Chux's Notebook

This is a collection of my notes which I refer to on a regular basis. Hope it is also helpful for others stumbling by.

I think of these notes as mountaineering pegs. Often at the time of studying a particular paper or topic, the concepts are clear. Over time, however, only a vague impression remains, and I can no longer tussle with the issues. So these notes serve as pegs that I hope to use to re-scale an old hill, or at least scale it with less effort than starting from scratch.

Generally, I highlight main points like so, and put in-line code or numbers like so.

Built with mdBook.

About Me

I am working as a Lead Data Scientist at GovTech Singapore, on a team called JumpStart. We power search and recommender systems for Singapore Government products like the job searching portal MyCareersFuture or the CareersFinder tool to discover next steps for one's career.

My LinkedIn profile is here.

Current Focus

2025 Goals

Incorporate semantic IDs into a generalizable search and recommender system:

Can function well with just a system prompt (guiding the overall goal of the system) and an item catalogue
Handles a changing catalogue gracefully
Good search and recommendation latency
Cheap (runs on a T4 or equivalent)

Research Questions

Optimizing LLM explanations based on implicit feedback
- How to optimize an LLM to provide better recommendation explanations by fine-tuning on implicit feedback?
Replacing BM25
- How to design a search system that matches BM25 performance at cold start and gradually improves with more data, without dropping below BM25 performance?
Precise Retrieval
- The common two tower approach to embedding retrieval leaves much to be desired
  - There is no natural score threshold at which items are deemed irrelevant. Traditionally, classifiers have a 0.5 score cut-off.
  - Embedding retrieval tends to retrieve unrelated items. This is a well documented problem. For example, Nike shoes retrieves Adidas Shoes.
  - Can we have embedding models that approximate AND / OR conditions that more naturally fit into the retrieval paradigm?
Efficient learning of semantic IDs
- LLMs can learn Semantic IDs as part of their language and recommend and reason about items once they learn the "language" using Supervised Fine Tuning
- But the danger is catastrophic forgetting and losing capabilities as they get fine tuned in this way
- Is there a more "natural" way for LLMs to learn such semantic IDs?
- How do we handle a changing catalogue in real-time gracefully?

Thoughts

Short notes and ideas.

On ML Experiments

On using LLMs (e.g. Claude Code) to assist with ML experiments, especially on replicating existing results from papers.

There is much value in replicating ML experiment results from papers.

Forces one to understand and implement the methods in the paper, no shortcuts
Helps one understand how sensitive the results are to hyperparameters, data settings etc.
Allows one to test out new methods or variations on the same task for an apples-to-apples comparison

Traditionally, this takes some time for each paper (days to weeks) depending on the complexity of the method, the compute cost, the expertise of the researcher who is performing the replication etc. There is also significant ambiguity in what defines a successful replication, since the paper may have inadvertently left out some implementation details, and the results might not be exactly replicable.

LLMs can make this process a lot faster, by:

Performing requisite data cleaning
Implementing the methods
Monitoring the training and validation runs and fixing bugs along the way

This raises the question: what are the respective roles of the LLM and researcher in such a process? From my experience, I have tested out these loosely described paradigms:

Iterative LLM-assisted coding
- Basically not too different from development before LLMs - the researcher codes from scratch and gets LLM to assist with each function, testing small pieces as we go along
Spec-driven LLM development + post-hoc checking
- Researcher starts with detailed specifications file, detailing as much as possible about the data, method, hyperparameters, implementation process, evaluation criteria etc.
- LLM implements semi-autonomously, checking against the specifications
- After each phase, researcher reviews the code to understand it, refactor it, simplify it etc.
Fully autonomous LLM
- Researcher provides the full paper to the LLM and high level specifications
- LLM implements end to end and submits results

I find that approach 1 produces the cleanest repo and understanding for the researcher. But it is also the slowest - significant time is spent debugging and doing the code edit -> run code loop manually. It's probably the approach that most researchers are using currently, to varying degrees of autonomy given to the LLM.

I could not get approach 3 to work. I tried running Claude Code in --dangerously-allow-permissions mode on runpod, and give it almost full autonomy in implementing and running the code. Some frustrating points:

LLM kept using synthetic data even for the full run, even when explicitly spelled out that synthetic data is only for dev runs. When asked why it was using synthetic data, it said it failed to download the data. But managed to download the data on the very next instruction.
LLM has no common sense. The paper detailed its method of generating synthetic queries on movielens using an LLM (since movielens only has interaction data) to train a search biencoder. Claude decided to generate synthetic queries of the form "What movie is ?", despite being told to use an LLM to generate from the movie text. Such hard-coded queries obviously led to very high accuracies of <code>> 99%</code>, but Claude did not know to terminate and investigate.</li> <li>LLM did not create meaningful evals. It mostly monitored train and validation loss, and even these were not very meaningful (it decided to sample only <code>500</code> points for validation). The researcher needs to be explicit about specifying these.</li> </ul> Some reflections on testing out approach 3: <ul> <li>Verification is hard in ML tasks. I tried providing the accuracy results from the paper to the LLM, asking it to benchmark against those figures to see if it is on the right track. The LLM was way off but did not know how to troubleshoot such problems. Normal test driven development also does not work for ML Tasks.</li> <li>Difficult to debug in fully autonomous setting. ML tasks are very sensitive to how the data is processed, hyperparameters etc. Even a small decision like left or right padding can significantly change results. Without diving into the code, I had no means of forming hypotheses about what was messing up the performance. LLMs are not well trained in this task of playing detective, and interrogating it did not yield much fruit.</li> <li>Unpleasant to understand code. Giving LLM full autonomy produces a lot of code that is intimidating to try to grok. I can probably do it with LLM assistance, but it produces the same daunting feeling as reviewing a large pull request. It also feels like unnecessary work as it entails spending time trying to understand lots of glue code.</li> </ul> Which leads me to approach 2. This is the approach I am currently testing out. The hope is that spec-driven development is a good middle ground. The researcher gets to use the specification to determine how far we want the LLM to sprint in each run, before we apply manual verification and understanding. This minimizes the chance of LLMs running completely off point and also makes it more palatable to understand a small piece of code each time. Front-loading the development of specifications also forces the researcher to adopt systems-level thinking and design for the whole experiment. This hopefully helps us move up a level of abstraction and move along faster. The main difficulties of approach 2 a-priori to me are: <ul> <li>Designing the right places to pause for manual verification and understanding</li> <li>Getting the LLM to follow our specs closely and stop at the correct checkpoints</li> </ul> Will test it out and see how it goes! <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="recent-trends-in-search--recommendations"><a class="header" href="#recent-trends-in-search--recommendations">Recent Trends in Search & Recommendations</a></h1> This document summarizes some important ideas in both recommender systems and search systems that have emerged in the past few years. It aims to be a helpful read for someone trying to get acquainted with modern practices in these fields. What we will cover: <ul> <li><a href="recsys/recent_trends.html#embedding-learning">Embedding learning</a>. Embedding learning for items and users is an upstream task that are used for downstream retrieval and ranking.</li> <li>Recommendations</li> <li>Search retrieval and ranking</li> </ul> <h2 id="embedding-learning"><a class="header" href="#embedding-learning">Embedding Learning</a></h2> Embedding learning is a foundational tool that powers all recommendation and search systems. Typically, users and items are represented by one or more embeddings that get fed into neural networks for prediction and recommendation tasks. Whilst many off-the-shelf LLM embedding models exist today, the best performing embeddings are often context-specific. For example, from the perspective of a job-seeker, we may want the embeddings for <code>Physical Education Teacher</code> and <code>Gym Trainer</code> to be similar, as the focus is on similar job functions / skills. However, from the perspective of a HR personnel trying to do manpower planning, we may want <code>Physical Education Teacher</code> to be more similar to <code>Executive at Singapore Sports School</code>, since the focus is more on eligible pathways for internal rotation (btw, I made this example up). Hence, there is a need to effectively learn embedding models in a specific product setting. Lee 2020 describes the task like so: <blockquote> (the goal is to) Learn embeddings that preserve the item-to-item similarities in a specific product context. </blockquote> It is common practice to have an upstream pipeline to learn item embeddings that get re-used for multiple downstream tasks (let's call this a "universal" embedding). For example, YouTube uses a common set of video embeddings used by all its models [Lee 2020]. To adapt the embeddings to a specific task, the task model itself can always add a simple translation MLP layer to translate the universal embedding so that the embeddings can be more performant in each task setting. Another way is to have each retrieval and ranker model learn its own embeddings for its specific task. <ul> <li>However, this leads to redundancy to store and update each set of embeddings, often in some in-memory feature store, which is costly.</li> <li>This approach is also often not as performant as learning a universal set of embeddings and then adapting the embeddings to each task via a small learned MLP layer</li> <li>Training from a frozen set of universal embeddings is also much more compute efficient, as we do not need to forward/backward propagate into a potentially large language model that generates the embeddings from text or some other modality.</li> </ul> We can think of upstream training of an embedding model as analogous to the common practice of pre-training large language models in a semi-supervised manner on a large text corpus. The LLM is then fine-tuned for specific applications. Although the performance of using multiple embeddings to represent an item has been proven to be more performant than using a single embedding (e.g. using ColBert), most companies seem to adopt a single embedding to reduce engineering complexity. For example, Pinterest briefly experimented with multiple embeddings per item but reverted to a single embedding to reduce costs. There are 3 primary ways of embedding representation in the literature: <ul> <li>ID-Based: This method maps an item's ID directly to an embedding. <ul> <li>f(item_id)→embedding</li> <li>This is often used in a task-specific model, but not usually used for a universal embedding model shared by downstream tasks, as we cannot represent new items that have no user activity</li> </ul> </li> <li>Content-Based: This is the most common approach, generating embeddings from an item's content, such as its text, images, or audio <ul> <li>f(item_text,item_image,...)→embedding</li> <li>Textual content is usually passed into an LLM to generate embeddings</li> <li>For other modalities, task-agnostic preprocessing may be employed. For example, Lee 2020 sampling frames from a video and runs them through a ResNET model to create raw embeddings before they are fed into the main embedder model</li> </ul> </li> <li>Graph-Based: This method can be seen as an extension of the content-based approach. It generates an embedding from the attributes of an item's neighbors in a network. <ul> <li>f(attributes of neighbours of item)→embedding</li> <li>This is particularly effective in social network products where network interactions are crucial, such as on LinkedIn, Pinterest, or Facebook</li> <li>While powerful, this method is more computationally expensive at both training and inference times because it requires access to the graph to compute embeddings</li> </ul> </li> </ul> Embeddings are typically trained in a contrastive learning manner, which simply means that we want to encourage related items (e.g. videos frequently co-watched) to have high similarity scores and unrelated items (e.g. randomly sampled vidoe pairs) to have low scores. Thus, we can think of contrastive learning as comprising 3 main components: <ul> <li>Positive Sampling. How we mine for related items from data is a non-trivial, but an oft-neglected topic. Usually some behavioural statistics are used, e.g. Lee 2020 uses <code>videos frequently co-watched by users</code> to determine positive pairs. The guiding principle is to err on the side of being stricter in the selection of positive samples to minimize the appearance of false positives in the data (which are more harmful than false negatives). <ul> <li>In JumpStart, we have explored using Pointwise Mutual Information score to discard item pairs with low PMI score. Recall that the PMI score is computed as: pmi(itema,itemb)=log2p(itema)⋅p(itemb)p(itema co-occuring with itemb)</li> <li>Considering co-occurrence statistics focuses only on the numerator. Normalizing by the marginal probabilities of the items reduces the incidence of popular items being labelled as related when they are actually not</li> <li>In my experience, the bottom <code>10%</code> of co-interacted items by pmi score can be safely discarded</li> <li>Another subtle issue is that power users contribute disproportionately to positive pairs. <ul> <li>If a given user interacts with n items, then he/she generates (2n)=2n(n−1) item co-occurrence pairs</li> <li>The quadratic scaling in n means that the item interactions from power users may dominate that of all other users, but these are also the least informative item pairs :'( as power users may have less specific item preferences</li> <li>Hence it is advisable to discard power users beyond a particular percentile to improve positive sampling</li> </ul> </li> </ul> </li> <li>Negative Sampling. Much of the research literature focuses on this issue. The reason is that as research data (e.g. <code>Movielens</code>) is fixed, the positive pairs are usually pre-determined in the data. But the negative samples are usually not pre-specified in the data, so methods are devised to sample: <ul> <li>Random sampling. This is the simplest and most common approach: for a given anchor item, we sample negatives uniformly from the item catalogue, omitting the positive item(s) for the anchor</li> <li>Impressed but not positive. This is also a simple approach if we have access to item impression data. The impressed items are often hard negatives because they were deemed to be relevant by the existing recommendation system and surfaced to the user. <ul> <li>The potential danger is that the model may only learn to distinguish negatives in a specific retrieval setting and is not robust to changes in the retrieval model</li> </ul> </li> <li>Hard negative mining. Various other methods exist to mine for hard negatives, e.g. using BM25 score or embedding similarity. We discuss two ideas on this below in a bit more detail.</li> <li>Usually, it is best to mix negatives from all of the above so that the model learns to distinguish positives from a wide range of negatives</li> </ul> </li> <li>Loss. Again, much ink has been spilt on this topic, but there are really only two main losses that we need to know: <ul> <li>Triplet loss.</li> <li>Cross entropy (or InfoNCE) loss.</li> </ul> </li> </ul> WORK IN PROGRESS To improve training, two ideas are critical: <ol> <li>Semi-hard Negative Mining: Originating from FaceNET, this technique addresses the problem of "false negatives" (items that are incorrectly labeled as negative) which can confuse the model. It focuses on training examples that are in a "goldilocks zone"—not too easy and not too hard. Lee (2020) demonstrated that mining for the hardest semi-hard negatives within a mini-batch is a computationally cheap and critical step for model performance.</li> <li>Smart Negative Sampling: Since random sampling often fails to find challenging negative examples, methods are devised to mine for them. This can involve using impression data, surrogate algorithms like BM25, or approximate-nearest-neighbor search. Lee (2020) also introduced using hierarchical clustering on a relational graph to find negatives from nearby clusters, which can be combined with semi-hard negative mining to prevent issues with false negatives.</li> </ol> <h3 id="from-matrix-factorization-to-sequential-recommendations"><a class="header" href="#from-matrix-factorization-to-sequential-recommendations">From Matrix Factorization to Sequential Recommendations</a></h3> Traditional matrix factorization methods learn a single, static user representation to predict their entire interaction history. This is problematic for users with diverse or evolving interests, as the model learns a "diluted" representation. While time-decayed ratings can help by giving more weight to recent interactions, the model still fails to learn the sequential nature of those interactions fully. Sequential recommenders address this by mirroring the actual recommendation task: given a user's history up to time t, predict their next action at time t+1. The user's representation is therefore dynamic and based on their recent actions. There are three main approaches to this: <ul> <li>Markov Approach: A simple but often effective method that uses the last interacted item as the primary feature to predict the next one.</li> <li>Pointwise Approach: This method involves "rolling back" features to what was available at the time of each action in the training data. While this explicitly computes every user-item interaction, it requires significant data-loading effort and careful feature engineering. Google chose this approach for its simplicity in distributed model training.</li> <li>Sequential Modelling Approach: This is analogous to language modeling, where the model learns to predict the next item in a sequence, learning from many items simultaneously.</li> </ul> <h3 id="case-studies-in-sequential-recommendations"><a class="header" href="#case-studies-in-sequential-recommendations">Case Studies in Sequential Recommendations</a></h3> <h4 id="covington-2016---pointwise-approach-at-youtube"><a class="header" href="#covington-2016---pointwise-approach-at-youtube">Covington (2016) - Pointwise Approach at YouTube:</a></h4> YouTube's influential 2016 paper detailed a pointwise approach for both retrieval and ranking. <ul> <li>Retrieval: The model used features like the average embedding of the last 50 video watches and search queries, along with user features like gender. A small MLP was trained with a sampled softmax loss to predict the next video watch. A crucial finding was that predicting the next watch was far more effective in online A/B tests than predicting a randomly held-out watch from the user's history.</li> </ul> [Diagram: Pointwise Approach for Retrieval - Covington 2016 architecture from slide 13] <ul> <li>"Example Age" Feature: To capture the time-sensitive nature of video popularity, an "Example Age" feature was introduced. This simple feature encodes when an action took place, allowing the model to learn the "popularity lifecycle" of items, where popularity often spikes shortly after upload and then fades.</li> </ul> [Diagram: Chart showing Class Probability vs. Days Since Upload with and without the "Example Age" feature, from slide 15] <ul> <li>Ranking: The ranking model was similar to the retrieval model but incorporated many more features, particularly those related to user-item historical interactions (e.g., "how many videos from this channel has the user watched?"). The paper noted that even simple transformations of features, like square roots, improved performance, highlighting the reality of manual feature engineering.</li> </ul> [Diagram: Pointwise Approach for Ranking - Covington 2016 architecture from slide 16] <h4 id="sasrec-2018---the-transformer-based-approach"><a class="header" href="#sasrec-2018---the-transformer-based-approach">SASRec (2018) - The Transformer-Based Approach:</a></h4> SASRec (Self-Attentive Sequential Recommendation) applied the transformer architecture to recommendations, treating a user's interaction history as a sequence of items to be predicted, much like words in a sentence. <ul> <li>Architecture: Using causal masking, the transformer embedding at each time step encodes all previous interactions. The training loss is typically a binary cross-entropy loss comparing the dot product of the model's prediction with the target item embedding against a random negative sample. This approach significantly outperforms traditional matrix factorization methods.</li> <li>Position Embeddings: Since transformers have no inherent notion of sequence, SASRec adds a learned position embedding to each item in the input sequence. This helps the model weigh recent items more heavily, as shown by visualizations of the attention matrix.</li> <li>Efficiency and Power: This paradigm is more efficient than the pointwise approach because it learns from N time steps simultaneously for each user. It also reduces the need for manual feature engineering, as the model can implicitly learn user-item interaction features.</li> </ul> [Diagram: Visualization of Average Attention Matrix from SASRec, with and without Positional Embeddings, from slide 21] A subsequent paper, BERT4Rec (2019), proposed a "masked token prediction" task similar to BERT, allowing the model to use future items as input. While it initially appeared superior, Klenitskiy (2023) showed this was due to a difference in loss functions; when both use the same sampled softmax loss, SASRec consistently performs better and trains faster. The conclusion is that SASRec with sampled softmax loss is the current industry standard. <h4 id="pinnerformer-2022---long-term-sequential-modeling-at-pinterest"><a class="header" href="#pinnerformer-2022---long-term-sequential-modeling-at-pinterest">PinnerFormer (2022) - Long-Term Sequential Modeling at Pinterest:</a></h4> PinnerFormer evolved from the need to switch from real-time user embedding computation to less costly daily batch jobs. <ul> <li>Long-Term Loss: The key idea is that at each time step, instead of only predicting the very next item (like SASRec), the model predicts a random item from the user's future interactions over a longer window (e.g., 28 days). This creates more stable embeddings and surprisingly, PinnerFormer beats SASRec even when retrained in real-time.</li> <li>Advanced Features: PinnerFormer uses pre-computed PinSage (graph-based) embeddings as a base. It distinguishes between different action types (e.g., "Pin Save" vs. "Long Click") by concatenating a learned "action" embedding to the item embedding. It also heavily relies on time features, using Time2Vec to encode timestamps, noting a significant performance drop without them.</li> <li>Training Details: The model uses a combination of in-batch and random negatives with Log-Q correction, and ensures each user in a mini-batch is weighted equally to avoid bias from users with very long histories.</li> </ul> [Diagram: Illustration of PinnerFormer's "Long Term Loss" compared to SASRec from slide 25] <h3 id="multi-task-learning-in-recommendations"><a class="header" href="#multi-task-learning-in-recommendations">Multi-Task Learning in Recommendations</a></h3> Recommendation systems often need to optimize for multiple objectives (e.g., clicks, saves, purchases). Several architectures address this, primarily for the pointwise approach: <ul> <li>ESMM (2018): This model explicitly encodes the relationship that a conversion can only happen after a click (<code>P(convert) = P(click) * P(convert|click)</code>). This allows the conversion model to benefit from the abundant click data while learning from the sparse conversion data.</li> </ul> [Diagram: ESMM architecture from slide 29] <ul> <li>Shared-Bottom Architecture: A common approach where bottom layers of a neural network are shared across tasks, with separate "towers" or heads for each specific task prediction. The final recommendation is a manually tuned weighted average of the logits from each head.</li> </ul> [Diagram: Shared-Bottom Model architecture from slide 30] <ul> <li>Mixture of Experts (MoE): An improvement over the shared-bottom model, MoE uses specialized "expert" layers. For each task, a gating network learns a weighted average over the outputs of these experts, allowing the model to modularize and handle potentially conflicting tasks more effectively.</li> </ul> [Diagram: Multi-gate Mixture-of-Expert (MoE) model architecture from slide 31] <ul> <li>PinnerFormer's Approach: For sequential models, PinnerFormer took a simpler route, finding that treating all positive actions (repins, clicks, etc.) as equal signals was the best all-around strategy for their use case.</li> </ul> <div class="table-wrapper"><table><thead><tr><th style="text-align: left">Training Objective</th><th style="text-align: left">10s Closeup</th><th style="text-align: left">10s Click</th><th style="text-align: left">Repin</th><th style="text-align: left">All</th></tr></thead><tbody> <tr><td style="text-align: left">10s Closeup</td><td style="text-align: left">0.27</td><td style="text-align: left">0.02</td><td style="text-align: left">0.09</td><td style="text-align: left">0.17</td></tr> <tr><td style="text-align: left">10s Click</td><td style="text-align: left">0.01</td><td style="text-align: left">0.49</td><td style="text-align: left">0.01</td><td style="text-align: left">0.12</td></tr> <tr><td style="text-align: left">Repin</td><td style="text-align: left">0.15</td><td style="text-align: left">0.03</td><td style="text-align: left">0.17</td><td style="text-align: left">0.13</td></tr> <tr><td style="text-align: left">Multi-task</td><td style="text-align: left">0.23</td><td style="text-align: left">0.28</td><td style="text-align: left">0.13</td><td style="text-align: left">0.23</td></tr> </tbody></table> </div> <hr /> <h2 id="ii-trends-in-search-systems"><a class="header" href="#ii-trends-in-search-systems">II. Trends in Search Systems</a></h2> While related to recommendations, search systems have distinct challenges and characteristics in practice. <h3 id="search-vs-recommendation"><a class="header" href="#search-vs-recommendation">Search vs. Recommendation</a></h3> <ul> <li>Relevance is Stricter: In search, there are objectively correct and incorrect results. Irrelevant results must not be surfaced, as they erode user trust.</li> <li>Latency is Critical: Users expect instant search results, making pre-computation of recommendations for a given query impossible.</li> <li>LLMs Have a Larger Impact: Search benefits more directly from advances in LLMs due to its text-centric nature and the overlap with fields like Information Retrieval and Question-Answering.</li> </ul> Search tasks exist on a spectrum from pure relevance (like academic Q&A) to pure personalization (which is essentially recommendation). E-commerce search lies in the middle, where relevance is key, but personalization plays a major role in user satisfaction given the large number of potentially relevant items for a broad query. <h3 id="academic-vs-industry-search"><a class="header" href="#academic-vs-industry-search">Academic vs. Industry Search</a></h3> Academic research often focuses on Question-Answering (Q&A) datasets like MSMarco or TriviaQA, where queries are well-formed questions with a single correct answer. In this setting, pre-trained LLMs perform exceptionally well, and it is easier to beat traditional baselines like BM25. <h4 id="standard-academic-models"><a class="header" href="#standard-academic-models">Standard Academic Models:</a></h4> <ul> <li>Ranking (Cross-Encoders): The standard is a cross-encoder architecture (often called MonoBERT), where a pre-trained language model like BERT takes a <code>[Query, Document]</code> pair as input and outputs a relevance score. Later papers showed that using a sampled softmax loss with many random negatives is more effective than the original binary cross-entropy loss. As models scale, so does performance, with RankLlama showing stronger results than older BERT or T5-based rankers. A 2023 paper by Sun even showed that zero-shot prompting with ChatGPT, using a "permutation generation" method and sliding windows, can achieve state-of-the-art performance on Q&A tasks without any fine-tuning.</li> </ul> [Diagram: Standard Cross Encoder / MonoBERT architecture from slide 39] <div class="table-wrapper"><table><thead><tr><th style="text-align: left">Ranking Model</th><th style="text-align: left">MRR@10 for MSMarco</th></tr></thead><tbody> <tr><td style="text-align: left">monoBERT</td><td style="text-align: left">37.2</td></tr> <tr><td style="text-align: left">RankT5</td><td style="text-align: left">43.4</td></tr> <tr><td style="text-align: left">LLaMA2-8b</td><td style="text-align: left">44.9</td></tr> </tbody></table> </div> <ul> <li>Retrieval (Two-Tower Models): The standard is Dense Passage Retrieval (DPR), which uses two separate encoders—one for the query and one for the document (passage). The model is trained with a sampled softmax loss to make the cosine similarity between a query and its positive document high, relative to negative documents. These negatives often include "in-batch" negatives (other positive documents in the same training batch) and "hard" negatives mined using BM25. To improve these retrievers, a common technique is distillation, where a powerful but slow cross-encoder (the "teacher") is used to train a faster two-tower model (the "student").</li> </ul> <div class="table-wrapper"><table><thead><tr><th style="text-align: left">Retrieval Model</th><th style="text-align: left">MRR@10 for MSMarco</th></tr></thead><tbody> <tr><td style="text-align: left">BM25</td><td style="text-align: left">18.4</td></tr> <tr><td style="text-align: left">ANCE (BERT-based)</td><td style="text-align: left">33.0</td></tr> <tr><td style="text-align: left">LLaMA2-8b</td><td style="text-align: left">41.2</td></tr> </tbody></table> </div> [Diagram: Illustration of a Mini Batch for Dense Passage Retrieval from slide 42] <h3 id="complexities-of-industry-search-systems"><a class="header" href="#complexities-of-industry-search-systems">Complexities of Industry Search Systems</a></h3> Industry search systems are significantly more complex, needing to balance relevance, latency, and personalization. <h4 id="case-study-pinterest-search-2024"><a class="header" href="#case-study-pinterest-search-2024">Case Study: Pinterest Search (2024)</a></h4> <ul> <li>Teacher-Student Distillation: Pinterest fine-tunes a large teacher LLM (Llama-3-8B) on 300k human-labeled query-pin pairs. To improve the teacher's performance, the text representation for each pin is enriched with its title, description, AI-generated image captions, high-engagement historical queries, and common board titles it's saved to.</li> </ul> <div class="table-wrapper"><table><thead><tr><th style="text-align: left">Model</th><th style="text-align: left">Accuracy</th><th style="text-align: left">AUROC 3+/4+/5+</th></tr></thead><tbody> <tr><td style="text-align: left">SearchSAGE</td><td style="text-align: left">0.503</td><td style="text-align: left">0.878/0.845/0.826</td></tr> <tr><td style="text-align: left">mBERT_base</td><td style="text-align: left">0.535</td><td style="text-align: left">0.887/0.864/0.861</td></tr> <tr><td style="text-align: left">T5_base</td><td style="text-align: left">0.569</td><td style="text-align: left">0.909/0.884/0.886</td></tr> <tr><td style="text-align: left">mDeBERTaV3_base</td><td style="text-align: left">0.580</td><td style="text-align: left">0.917/0.892/0.895</td></tr> <tr><td style="text-align: left">XLM-ROBERTalarge</td><td style="text-align: left">0.588</td><td style="text-align: left">0.919/0.897/0.900</td></tr> <tr><td style="text-align: left">Llama-3-8B</td><td style="text-align: left">0.602</td><td style="text-align: left">0.930/0.904/0.908</td></tr> </tbody></table> </div> <ul> <li>Scaling Up with Distilled Data: This powerful teacher model then generates labels for 30 million un-labeled data points. A small, fast student MLP model is then trained on this massive distilled dataset. This student model recovers up to 91% of the teacher's accuracy and, because the teacher was multilingual, the student model also learns to handle multiple languages effectively despite being trained only on English human-annotated data.</li> </ul> <div class="table-wrapper"><table><thead><tr><th style="text-align: left">Training Data</th><th style="text-align: left">Accuracy</th><th style="text-align: left">AUROC 3+/4+/5+</th></tr></thead><tbody> <tr><td style="text-align: left">0.3M human labels</td><td style="text-align: left">0.484</td><td style="text-align: left">0.850/0.817/0.794</td></tr> <tr><td style="text-align: left">6M distilled labels</td><td style="text-align: left">0.535</td><td style="text-align: left">0.897/0.850/0.841</td></tr> <tr><td style="text-align: left">12M distilled labels</td><td style="text-align: left">0.539</td><td style="text-align: left">0.903/0.856/0.847</td></tr> <tr><td style="text-align: left">30M distilled labels</td><td style="text-align: left">0.548</td><td style="text-align: left">0.908/0.860/0.850</td></tr> </tbody></table> </div> [Diagram: Teacher-Student distillation process at Pinterest from slide 47] <h4 id="case-study-baidu-search-2021-2023"><a class="header" href="#case-study-baidu-search-2021-2023">Case Study: Baidu Search (2021, 2023)</a></h4> <ul> <li>Bootstrapping Labels with Weak Signals: Instead of training a large teacher model, Baidu uses a simple decision tree model trained on "weak relevance signals" (e.g., click-through rates, long-click rates, skip rates) to predict a "calibrated relevance score" for unlabeled data. Training a cross-encoder on this calibrated score was shown to be far more effective than training on raw clicks alone.</li> </ul> [Diagram: Baidu's decision tree model over weak relevance signals from slide 48] <ul> <li>Query-Sensitive Summary: To avoid the latency of processing full documents, Baidu developed an algorithm to extract the most relevant sentences from a document based on word-match scores with the query. This summary is then used as input for the ranker, significantly boosting offline performance.</li> <li>Modular Attention for Speed: To accelerate their cross-encoder, Baidu uses a modular attention mechanism. Self-attention is first applied independently within the <code>[query-title]</code> segment and the <code>document summary</code> segment for several layers, before a few final layers of full cross-attention are applied. This structure reduces computation and increases inference speed by 30%.</li> </ul> [Diagram: Baidu's Modular Attention architecture from slide 50] <ul> <li>From Relevance to Satisfaction: A 2023 paper from Baidu argued that relevance alone is insufficient and that models must optimize for user satisfaction. They engineered features for quality (e.g., number of ads), authority (e.g., PageRank), and recency, and even built a query analyzer to determine if a query was authority or recency-sensitive. These features, along with relevance, were used to train a model that boosted performance by 2.5% over a pure relevance model. Notably, numeric satisfaction features were simply normalized and appended directly into the text input for the LLM to process.</li> </ul> <h4 id="case-study-google-dcnv2-2020"><a class="header" href="#case-study-google-dcnv2-2020">Case Study: Google DCNv2 (2020)</a></h4> Google's Deep & Cross Network V2 (DCNv2) addresses a weakness in standard MLP rankers: they are not ideal for explicitly modeling feature interactions (e.g., <code>user_is_interested_in_topic_X</code> AND <code>item_is_about_topic_X</code>). DCNv2 introduces explicit "cross layers" that are designed to generate these feature crosses efficiently before feeding them into a standard deep network. This architecture provides performance gains over a standard MLP with the same number of parameters. [Diagram: Google DCNv2 architecture with Cross Layers from slide 53] <h4 id="case-studies-embedding-based-retrieval-ebr-at-facebook-and-taobao"><a class="header" href="#case-studies-embedding-based-retrieval-ebr-at-facebook-and-taobao">Case Studies: Embedding-Based Retrieval (EBR) at Facebook and Taobao</a></h4> EBR systems are used for semantic search but can struggle with relevance compared to traditional keyword-based lexical matching. <ul> <li>Facebook's Hybrid Approach: Facebook integrates nearest-neighbor search directly into its in-house lexical search engine. A query can contain both standard lexical terms (e.g., <code>location:seattle</code>) and an <code>nn</code> operator that finds items within a certain embedding-distance radius. This ensures documents fulfill both lexical and semantic requirements. Facebook's query and document encoder towers use a mix of term-based (char-3-grams) and LLM-based representations.</li> <li>Taobao's Relevance Filters: Taobao found their EBR system sometimes returned irrelevant but semantically close items (e.g., a search for "Nike shoes" returning "Adidas shoes"). Their solution was to apply explicit lexical boolean filters on top of the ANN search (e.g., <code>ANN search + Brand:Nike + Product:Shoes</code>). While this sacrifices some of the "fuzzy" semantic search capability, it guarantees relevance. Taobao also found that optimizing for engagement (clicks) does not guarantee relevance, and ultimately chose to accept slightly lower engagement to ensure a higher percentage of relevant results. Their query representation is enhanced by using cross-attention between the query and the user's historical interactions to personalize the query embedding.</li> </ul> <div class="table-wrapper"><table><thead><tr><th style="text-align: left">Experiment</th><th style="text-align: left">Engagement %</th><th style="text-align: left">Relevance %</th></tr></thead><tbody> <tr><td style="text-align: left">Baseline</td><td style="text-align: left">85.6%</td><td style="text-align: left">71.2%</td></tr> <tr><td style="text-align: left">Baseline + personalization</td><td style="text-align: left">86.4%</td><td style="text-align: left">71.4%</td></tr> <tr><td style="text-align: left">Baseline + Lower temperature</td><td style="text-align: left">85.5%</td><td style="text-align: left">79.0%</td></tr> <tr><td style="text-align: left">Baseline + all</td><td style="text-align: left">84.7%</td><td style="text-align: left">80.0%</td></tr> </tbody></table> </div><div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h2 id="lightgbm-memory"><a class="header" href="#lightgbm-memory">LightGBM Memory</a></h2> TLDR: Solutions for memory issues during training of a LightGBM model: <ol> <li>Cast numeric values into <code>np.float32</code> to save data space</li> <li>Keep <code>num_leaves <= 100</code> (or some reasonable number)</li> <li>If feature dimension is large (e.g. <code>M >= 1000</code>), try <code>colsample_bytree = 0.1</code>, although this might not help too much if the bottleneck is during bin histogram construction (rather than the actual training)</li> <li>If number of rows and features are both large (e.g. <code>N >= 1_000_000</code> and <code>M >= 1000</code>, i.e. <code>>= 4 GB</code>) then the data itself is taking up a lot of memory. It would be worthwhile to put the data on disk and use <code>lgb.Dataset</code> by providing the file path as the data argument instead. Then, we should set <code>two_round=True</code> for the train method params. The <a href="https://github.com/microsoft/LightGBM/issues/1138#issuecomment-353857567">explanation</a> for two round is rather unclear, but it should help with memory when <code>Dataset</code> is loading from disk (rather than from a <code>numpy.array</code> in memory). For this option, I had some trouble getting it to work with categorical columns.</li> </ol> For more details can refer to the experiments below. <h3 id="experiments"><a class="header" href="#experiments">Experiments</a></h3> I often run into memory issues running LightGBM. So here are some experiments to measure memory usage and understand how hyperparameters can affect memory usage. The function of interest is the <code>fit</code> method for the learn to rank task. <pre><code class="language-python">import lightGBM as lgb def f(): model = lgb.LGBMRanker(**params, objective="lambdarank") model.fit( X=data, y=y, group=groups, ) </code></pre> The memory usage is measured using the <code>memory_profiler</code> module, which checks the memory usage at .1 second intervals. The maximum is then taken to represent the maximum memory usage of the fit function. We also take note of the size of the data itself (using <code>data.nbytes</code>) and subtract that away to get closer to the LightGBM memory usage. Do note that this memory profiling is not very rigorous, so the results are best for relative comparison within each experiment rather than across experiments. <pre><code class="language-python">from memory_profiler import memory_usage def run(params): mem_usage = memory_usage(f) return max(mem_usage) / 1000 # GB </code></pre> We set the default parameters as follows and generate the data this way. For the experiments below, the default parameters are used unless specified otherwise. <pre><code class="language-python">DEFAULT_PARAMS = { "N": 200000, # number of instances "M": 500, # feature dimension "n_estimators": 100, "num_leaves": 100, "histogram_pool_size": -1, } data = np.random.randn(DEFAULT_PARAMS["N"], DEFAULT_PARAMS["M"]) groups = [20] * int(N / 20) # assume each session has 20 rows y = np.random.randint(2, size=N) # randomly choose 0 or 1 </code></pre> Large <code>num_leaves</code> can get very memory intensive. We should not need too many leaves, so generally using <code>num_leaves <= 100</code> and increasing the number of estimators seems sensible to Gme. <ul> <li>num_leaves: <code>10</code>, Maximum memory usage: 2.28 GB - 0.80 GB = <code>1.48 GB</code></li> <li>num_leaves: <code>100</code>, Maximum memory usage: 2.52 GB - 0.80 GB = <code>1.72 GB</code></li> <li>num_leaves: <code>1000</code>, Maximum memory usage: 4.04 GB - 0.80 GB = <code>3.24 GB</code></li> </ul> Increasing <code>n_estimators</code> doesn't seem to raise memory much, but increases run time because each tree is fitted sequentially on the residual errors, so it cannot be parallelized. <ul> <li>n_estimators: <code>10</code>, Maximum memory usage: 2.28 GB - 0.80 GB = <code>1.48 GB</code></li> <li>n_estimators: <code>100</code>, Maximum memory usage: 2.53 GB - 0.80 GB = <code>1.73 GB</code></li> <li>n_estimators: <code>1000</code>, Maximum memory usage: 2.69 GB - 0.80 GB = <code>1.89 GB</code></li> </ul> Increasing <code>N</code> increases memory sublinearly. It seems that the data size itself will be more of a problem than the increase in LightGBM memory usage as <code>N</code> increases. For extremely large <code>N</code>, we can also set the <code>subsample</code> parameter to use only a fraction of the training instances for each step (i.e. stochastic rather than full gradient descent). By default <code>subsample=1.0</code>. <ul> <li>N: <code>1,000</code>, Maximum memory usage: 0.38 GB - 0.00 GB = <code>0.38 GB</code></li> <li>N: <code>10,000</code>, Maximum memory usage: 0.45 GB - 0.04 GB = <code>0.41 GB</code></li> <li>N: <code>100,000</code>, Maximum memory usage: 1.46 GB - 0.40 GB = <code>1.06 GB</code></li> <li>N: <code>1,000,000</code>, Maximum memory usage: 6.12 GB - 4.00 GB = <code>2.12 GB</code></li> <li>N: <code>2,000,000</code>, Maximum memory usage: 10.48 GB - 8.00 GB = <code>2.48 GB</code></li> </ul> In contrast to <code>N</code>, memory usage is quite sensitive to <code>M</code>, seems to increase linearly when <code>M</code> gets large. <code>M=10,000</code> blows up my memory. I suppose this could be mitigated by setting <code>colsample_bytree</code> or <code>colsample_bynode</code> to sample a smaller subset. <ul> <li>M: <code>100</code>, Maximum memory usage: 2.08 GB - 0.16 GB = <code>1.92 GB</code></li> <li>M: <code>1000</code>, Maximum memory usage: 4.92 GB - 1.60 GB = <code>3.32 GB</code></li> <li>M: <code>2000</code>, Maximum memory usage: 9.69 GB - 3.20 GB = <code>6.49 GB</code></li> <li>M: <code>3000</code>, Maximum memory usage: 14.35 GB - 4.80 GB = <code>9.55 GB</code></li> </ul> To deal with the high memory usage of large <code>M</code>, we can set <code>colsample_bytree</code> which samples a subset of columns before training each tree. This will help to mitigate the memory usage. For this experiment, we set <code>M=2000</code> to simulate data with high number of dimensions. <ul> <li><code>colsample_bytree</code>: 0.1, Maximum memory usage: 8.60 GB - 3.20 GB = <code>5.40 GB</code></li> <li><code>colsample_bytree</code>: 0.2, Maximum memory usage: 9.58 GB - 3.20 GB = <code>6.38 GB</code></li> <li><code>colsample_bytree</code>: 0.4, Maximum memory usage: 10.06 GB - 3.20 GB = <code>6.86 GB</code></li> <li><code>colsample_bytree</code>: 0.6, Maximum memory usage: 10.07 GB - 3.20 GB = <code>6.87 GB</code></li> <li><code>colsample_bytree</code>: 0.8, Maximum memory usage: 10.46 GB - 3.20 GB = <code>7.26 GB</code></li> </ul> In contrast, setting <code>colsample_bynode</code> does not help memory usage at all. Not too sure why, but I suppose since multiple nodes for the same tree can be split at the same time, the full feature set still has to be kept in memory. <ul> <li><code>colsample_bynode</code>: 0.1, Maximum memory usage: 10.49 GB - 3.20 GB = <code>7.29 GB</code></li> <li><code>colsample_bynode</code>: 0.2, Maximum memory usage: 10.49 GB - 3.20 GB = <code>7.29 GB</code></li> <li><code>colsample_bynode</code>: 0.4, Maximum memory usage: 10.49 GB - 3.20 GB = <code>7.29 GB</code></li> <li><code>colsample_bynode</code>: 0.6, Maximum memory usage: 10.49 GB - 3.20 GB = <code>7.29 GB</code></li> <li><code>colsample_bynode</code>: 0.8, Maximum memory usage: 10.48 GB - 3.20 GB = <code>7.28 GB</code></li> </ul> Tweaking <code>boosting</code> and <code>data_sample_strategy</code> don't seem to affect memory usage too much. Using <code>dart</code> seems to require a bit more memory than the traditional <code>gbdt</code>. <ul> <li><code>data_sample_strategy</code>: bagging, <code>boosting</code>: gbdt, Maximum memory usage: 8.90 GB - 3.20 GB = <code>5.70 GB</code></li> <li><code>data_sample_strategy</code>: goss, <code>boosting</code>: gbdt, Maximum memory usage: 9.58 GB - 3.20 GB = <code>6.38 GB</code></li> <li><code>data_sample_strategy</code>: bagging, <code>boosting</code>: dart, Maximum memory usage: 9.81 GB - 3.20 GB = <code>6.61 GB</code></li> <li><code>data_sample_strategy</code>: goss, <code>boosting</code>: dart, Maximum memory usage: 9.80 GB - 3.20 GB = <code>6.60 GB</code></li> </ul> Another bottleneck we can tackle is to realize that LightGBM is a two-stage algorithm. In the first stage, LightGBM uses the full dataset to construct bins for each numeric variable (controlled by the <code>max_bins</code> argument) based on the optimal splits. In the second stage, these discretized bins are then used to map and split the numeric variables during the actual training process to contruct trees. From my understanding, the first stage cannot be chunked as it requires the full dataset, but the second stage can be chunked (as per any stochastic gradient descent algorithm) where a fraction of the dataset is loaded at each time. Hence, the real bottleneck appears to be the first stage, when the bins are constructed. According to this <a href="https://github.com/microsoft/LightGBM/issues/1146">thread</a>, we can separate the memory usage between the two stages by using <code>lgb.Dataset</code>. First, we initialize the Dataset object and make sure to set <code>free_raw_data=True</code> (this tells it to free the original data array after the binning is done). Then, we trigger the actual dataset construction using <code>dataset.construct()</code>. Thereafter, we are free to delete the original data array to free up memory for the actual training. The following code illustrates this concept. <pre><code class="language-python">dataset = lgb.Dataset(data=data, label=y, group=groups, free_raw_data=True) del data dataset.construct() lgb.train(params=params, train_set=dataset) </code></pre> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="tf-idf"><a class="header" href="#tf-idf">TF-IDF</a></h1> Term Frequency - Inverse Document Frequency is a well known method for representing a document as a bag of words. For a given corpus C, we compute the IDF value for each word w by taking idfw:=logdfw1, with dfw denoting the number of documents in C containing the word w. The document d is represented by a vector of length corresponding to the number of unique words in C. Each element of the vector will be a tf-idf value for the word, i.e. tfidfwd:=tfwd⋅idfw, where tfwd represents the term frequency of the word w in document d. Sometimes, we may l1 or l2 normalize the tf-idf vector so that the dot product between document vectors represents the cosine similarity between them. <h2 id="bayesian-smoothing"><a class="header" href="#bayesian-smoothing">Bayesian Smoothing</a></h2> We may want to apply some bayesian smoothing to the tfwd terms to avoid spurious matches. For example, suppose that a rare word wr appears only in documents d1 and d2 in the entire corpus just by random chance. The idfwr will be a large value, and hence documents d1 and d2 will have a high cosine similarity just because of this rare word. For the specific setting I am considering, we can deal with this problem using bayesian smoothing. The setting is as follows: <ul> <li>Each document d represents a job, and each job is tagged to an occupation o</li> <li>An occupation can have one or more jobs tagged to it</li> <li>We wish to represent each occupation as a TF-IDF vector of words</li> </ul> To apply bayesian smoothing to this scenario, notice that we only need to smooth the term frequencies tfwd. Since the IDF values idfw are estimated across the whole corpus, we can assume that those are relatively reliable. And since term frequencies are counts, we can use a poisson random variable to represent them. See <a href="https://www.y1zhou.com/series/bayesian-stat/bayesian-stat-bayesian-inference-poisson/">reference</a> for a primer on the gamma-poisson bayesian inference model. Specifically, we assume that θwo is the poisson parameter that dictates the term frequency of w in any document belonging to o, i.e. tfwd∼Poisson(θwo). We treat the observed term frequency tfwd for each document d belonging to o as a data sample to update our beliefs about θwo. We start with an uninformative gamma prior for θwo, and obtain the MAP estimate for θ^wo as below, with dfd∈o denoting the number of documents that belong to occupation o. θ^wo=b+dfd∈oa+∑d∈otfwd−1 We can thus use this formula to obtain posterior estimates for each θ^wo. One possible choice of the prior parameters a and b is to set a to be the mean term frequency for word w per document in the entire corpus, and to set b:=1. This prior corresponds to θ^wo following a gamma distribution with mean a and variance a, which seems to be a reasonable choice that can be overrided by a reasonable amount of data. The posterior variance, which may also be helpful in quantifying the confidence of this estimate, is: Var(θ^wo)=(b+dfd∈o)2a+∑d∈otfwd Finally, after obtaining the posterior estimates for each θ^wo, we can just use them as our term frequencies and multiply them by the IDF values as per normal. We can also apply l1 or l2 normalization thereafter to the tf-idf vectors. This method should produce tf-idf vectors that are more robust to the original problem of spurious matches. For illustration, for a very rare word wr, a will be a low value close to 0 (say 0.01). Suppose we were to observe n number of new documents, each containing one occurrence of word wr. Then the posterior estimate of θ^w will update as follows: <div class="table-wrapper"><table><thead><tr><th>n</th><th>θ^w</th></tr></thead><tbody> <tr><td>1</td><td>0.005</td></tr> <tr><td>2</td><td>0.337</td></tr> <tr><td>3</td><td>0.503</td></tr> <tr><td>4</td><td>0.602</td></tr> <tr><td>...</td><td>...</td></tr> <tr><td>20</td><td>0.905</td></tr> </tbody></table> </div> As desired, the estimate for θ^w starts off at a very small value and gradually approaches the true value 1. This will help mitigate the effect of spurious matches. If we desire for the update to match the data more quickly, we can simply scale a and b down by some factor, e.g. now a:=5a=0.002 and b:=5b=0.2. Then we have: <div class="table-wrapper"><table><thead><tr><th>n</th><th>θ^w</th></tr></thead><tbody> <tr><td>1</td><td>0.001</td></tr> <tr><td>2</td><td>0.455</td></tr> <tr><td>3</td><td>0.626</td></tr> <tr><td>4</td><td>0.715</td></tr> <tr><td>...</td><td>...</td></tr> <tr><td>20</td><td>0.941</td></tr> </tbody></table> </div> As a final detail, note that the update formula can result in negative estimates if a<1 and ∑dtfwd=0. The small negative value is probably not a big problem for our purposes, but we could also resolve it by setting the negative values to zero if desired. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="cross-encoders"><a class="header" href="#cross-encoders">Cross Encoders</a></h1> Cross encoder is a type of model architecture used for re-ranking a relatively small set of candidates (typically <code>1,000</code> or less) with great precision. In the Question-Answering or machine reading literature, typically the task involves finding the top matching <code>documents</code> to a given <code>query</code>. A typical task is the MS MARCO dataset, which seeks to find the top documents that are relevant to a given bing query. <h2 id="basic-setup"><a class="header" href="#basic-setup">Basic Setup</a></h2> Typically, the base model is some kind of pre-trained BERT model, and a classification head is added on top to output a probability. Each <code>(query, document)</code> pair is concatenated with <code>[SEP]</code> token in-between to form a sentence. The sentence is fed into the classification model to output a probability. The model is trained using binary cross-entropy loss against <code>0,1</code> labels (irrelevant or relevant). This is the setup used by <Nogeuira 2019>, possibly the first paper to propose the cross encoder. Some specifics for their setup: <ul> <li>The query is truncated to max <code>64 tokens</code>, while the passage is truncated such that the concatenated sentence is max <code>512 tokens</code>. They use the <code>[CLS]</code> embedding as input to a classifier head.</li> <li>The loss for a single query is formulated as below. si,sj refers to the score from the classifier model, Jpos refers to the documents that are relevant, and Jneg refers to documents in the top <code>1,000</code> retrieved by BM25 that are not relevant. Note that this results in a very imbalanced dataset.</li> </ul> L=−j∈Jpos∑log(sj)−j∈Jneg∑log(1−sj) <ul> <li>The model is fine-tuned with a batch size of 128 sentence pairs for 100k batches.</li> </ul> As opposed to bi-encoders (or dual encoders), which take a dot product between the <code>query</code> embedding and the <code>document</code> embedding, we cannot pre-compute embeddings in the cross encoder setting, because the cross encoder requires a forward pass on the concatenated <code>(query, document)</code> pair. Due to the bi-directional attention on the full concatenated sentence, we need the full sentence before we can compute the score, which requires the <code>query</code> that we only see at inference time. Hence, the cross encoder is limited to reranking a small set of candidates as it requires a full forward pass on each <code>query, candidate_document</code> pair separately. <h2 id="contrastive-loss"><a class="header" href="#contrastive-loss">Contrastive Loss</a></h2> The vanilla binary cross entropy loss proposed above may be thought of as a <point-wise> loss, in which each document is either relevant or irrelevant in absolute terms. However, treating relevance as a <relative> concept often better reflects reality. For example, given the first page of search results for a Google query, most of the documents should be relevant to some extent, but some are more relevant than the rest (and get clicked on). Simply treating all clicks as relevant and all non-clicks as irrelevant naively ignores the context (i.e. the neighbouring search results) in which the clicks were generated. It assumes that across query sessions, the average <absolute> level of relevance of the results is comparable. Treating relevance as a <relative> concept within the same query session weakens this assumption and hence often works better. Thus <Gao 2021> proposes the Local Contrastive Estimation loss. For a given query <code>q</code>, a positive document dq+ is selected, and a few negative documents dq− are sampled using a retriever (e.g. BM25). The contrastive loss then seeks to maximize the softmax probability of the positive document against the negative documents. LLCE=∣Q∣1q∈Q, Gq∑−log∑dq−∈Gqexp(sq, dq−)exp(sq, dq+) It is confirmed in multiple experiments in Gao 2021 and Pradeep 2022 that LCE consistently out-performs point-wise cross entropy loss. Furthermore, the performance consistently improves as the number of negative documents per query (i.e. ∣Gq∣) increases. In Gao 2021, up to <code>7</code> negatives (i.e. batch size of <code>8</code>) were used. Pradeep 2022 shows that increasing the batch size up to <code>32</code> continues to yield gains consistently (albeit diminishingly). <h2 id="other-details"><a class="header" href="#other-details">Other details</a></h2> Pradeep 2022's experiments show that using a stronger retrieval model (a ColBERT-based model) during inference generates slight gains in final performance (as opposed to BM25). Although Gao 2021 argues that it is also important to use the same retrieval model during model training (so that the cross encoder sees the same distribution of negatives during training and inference), Pradeep 2022 argues that the alignment is not as important as the stronger retrieval performance during inference. <h2 id="references"><a class="header" href="#references">References</a></h2> <ul> <li><a href="https://ar5iv.labs.arxiv.org/html/1901.04085">Nogueira 2019 - Passage Re-ranking with BERT</a></li> <li><a href="https://ar5iv.labs.arxiv.org/html/2101.08751">Gao 2021 - Rethink Training of BERT Rerankers in Multi-Stage Retrieval Pipeline</a></li> <li><a href="https://cs.uwaterloo.ca/~jimmylin/publications/Pradeep_etal_ECIR2022.pdf">Pradeep 2022 - A Bag of Tricks for Improving Cross Encoder Effectiveness</a></li> <li><a href="https://huggingface.co/blog/train-reranker">Huggingface reranker training</a></li> </ul> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="sentencetransformers"><a class="header" href="#sentencetransformers">SentenceTransformers</a></h1> SentenceTransformers is a useful library for training various BERT-based models, including two-tower embedding models and cross encoder reranking models. <h2 id="cross-encoder"><a class="header" href="#cross-encoder">Cross Encoder</a></h2> SentenceTransformers v4.0 updated their cross encoder training interface (see <a href="https://huggingface.co/blog/train-reranker">the v4.0 blogpost</a>). Here we try to follow the key components for cross encoder training using their API. The main class for training is <code>CrossEncoderTrainer</code>. We rely on a Huggingface <code>datasets.Dataset</code> class to provide training and validation data. <code>CrossEncoderTrainer</code> requires that the dataset format matches the chosen loss function. The <a href="https://sbert.net/docs/cross_encoder/loss_overview.html">loss overview page</a> provides a summary of cross encoder losses and the required dataset format. In general for cross encoder training, we have two sentences which are either positively or negatively related to each other. Which loss function we choose depends on the specific dataset format we possess. <h3 id="binarycrossentropyloss"><a class="header" href="#binarycrossentropyloss">BinaryCrossEntropyLoss</a></h3> Use this loss if we have inputs in the form of <code>(sentence_A, sentence_B)</code> and a label of either <code>0: negative, 1: positive</code> or a float score between <code>0-1</code>. In the huggingface <code>dataset</code>, we would need to ensure that the label column is named <code>label</code> or <code>score</code>, and have two other input columns corresponding to <code>sentence_A</code> and <code>sentence_B</code>. For <code>sentence_transformers</code> package in general, order of columns matter, so we should set it to <code>sentence_A, sentence_B, label</code>. Inspecting the <a href="https://github.com/UKPLab/sentence-transformers/blob/master/sentence_transformers/cross_encoder/losses/BinaryCrossEntropyLoss.py#L9-L116">source code</a> would show that each sentence pair is tokenized and encoded by the cross encoder model. The cross encoder must output a single logit (i.e. initialized with <code>num_labels=1</code>). Thus we get a prediction vector x of dim <code>batch_size</code>. The <code>torch.nn.BCEWithLogitsLoss</code> is then used to compute the binary cross entropy loss of the prediction logits against the actual labels y, according to the standard bce loss: L(xi,yi)=−wi[yi⋅logσ(xi)+(1−yi)⋅log(1−σ(xi))] This is a simple and effective loss. The user should ensure that the labels are well distributed (between 0 and 1) without any severe class imbalance. <h3 id="crossentropyloss"><a class="header" href="#crossentropyloss">CrossEntropyLoss</a></h3> The CrossEntropyLoss is used for a classification task, where for a given input sentence pair <code>(sentence_A, sentence_B)</code>, the label is a class. For example, we may have data where each sentence pair is tagged to a <code>1-5</code> rating scale. We need to instatiate the <code>CrossEncoder</code> class with <code>num_labels=num_classes</code> for this use case. This creates a prediction head for each class. Looking at the <a href="https://github.com/UKPLab/sentence-transformers/blob/master/sentence_transformers/cross_encoder/losses/CrossEntropyLoss.py#L8-L84">source code</a>, we see that this loss simply takes the prediction logits from the model (of dimension <code>num_labels</code>) and computes the <code>torch.nn.CrossEntropyLoss</code> against the actual labels. Note that the cross entropy loss takes the following form. Given <code>num_labels=C</code> and logits of x1,...,xC, where the correct label is index y, we have: L(x,y)=−log(∑j=1Cexjexy) <h3 id="multiplenegativesrankingloss"><a class="header" href="#multiplenegativesrankingloss">MultipleNegativesRankingLoss</a></h3> This is basically InfoNCE loss or in-batch negatives loss. The inputs to this loss can take the following forms: <ul> <li><code>(anchor, positive)</code> sentences</li> <li><code>(anchor, positive, negative)</code> sentences</li> <li><code>(anchor, positive, negative_1, ..., negative_n)</code> sentences</li> </ul> The <a href="https://sbert.net/docs/package_reference/cross_encoder/losses.html#multiplenegativesrankingloss">documentation page</a> has a nice description of what this loss does: Given an anchor, assign the highest similarity to the corresponding positive document out of every single positive and negative document in the batch. Diving into the <a href="https://github.com/UKPLab/sentence-transformers/blob/master/sentence_transformers/cross_encoder/losses/MultipleNegativesRankingLoss.py#L12-L191">source code</a>: <ul> <li>The inputs are <code>list[list[str]]</code>, where the outer list corresponds to <code>[anchor, positive, *negatives]</code>. The inner list corresponds to the batch size.</li> <li><code>scores</code> of dimension <code>(batch_size)</code> are computed for each anchor, positive pair</li> <li><code>get_in_batch_negatives</code> is then called to mine negatives for each anchor. <ul> <li>candidates (positive and negatives) are extracted at <code>inputs[1:]</code> and flattened into a long list</li> <li>A mask is created such that for each anchor, all the matching positive and negative candidates are masked out (not participating)</li> <li>The matching negatives do not participate because they will be added later on</li> <li>Amongst the remaining negatives, <code>torch.multinomial</code> is used to select <code>self.num_negatives</code> number of documents per anchor at random</li> <li><code>self.num_negatives</code> defaults to <code>4</code></li> <li>These randomly selected negative texts are then returned as <code>list[str]</code></li> </ul> </li> <li>For each negative in num_negatives mined in-batch negatives: <ul> <li><code>score</code> of dimension <code>(batch_size)</code> is computed for the anchor, negative pair</li> <li>The result is appended to <code>scores</code></li> </ul> </li> <li>Similarly, for each hard matching negative: <ul> <li><code>score</code> of dimension <code>(batch_size)</code> is computed for the anchor, hard negative pair</li> <li>The result is appended to <code>scores</code></li> </ul> </li> </ul> Now <code>scores</code> is passed into <code>calculate_loss</code>: <ul> <li>Recall that <code>scores</code> is a list of tensors where the outer list is size <code>1 + num_rand_negatives + num_hard_negatives</code>, and each tensor is of dimension <code>batch_size</code></li> <li>Thus <code>torch.cat</code> + <code>tranpose</code> is called to make it <code>(batch_size, 1 + num_rand_negatives + num_hard_negatives)</code></li> <li>Note that for each row, the first column corresponds to the positive document</li> <li>Hence the <code>labels</code> may be created as <code>torch.zeros(batch_size)</code></li> <li>Then <code>torch.nn.CrossEntropyLoss()(scores, labels)</code> may be called to get the loss</li> </ul> This sums up the loss computation for <code>MultipleNegativesRankingLoss</code>. <h3 id="cachedmultiplenegativesrankingloss"><a class="header" href="#cachedmultiplenegativesrankingloss">CachedMultipleNegativesRankingLoss</a></h3> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="collaborative-filtering"><a class="header" href="#collaborative-filtering">Collaborative Filtering</a></h1> Collaborative filtering is typically done with implicit feedback in the RecSys setting. In this setting, interactions are often very sparse. Most of the time, only positive signals are recorded, but a non-interaction could either mean (i) user dislikes the item or (ii) the user was not exposed to the item. Hence, we cannot use algorithms like SVD which assume no interactions as irrelevance. A useful repository is https://github.com/recommenders-team/recommenders. A generic and fairly common architecture for the collaborative filtering model is to embed each user and item into separate fixed size vectors, and use the cosine similarity between the vectors to represent a score. This score is fed into a cross entropy loss against the labelled relevance of user to item to train the embeddings. <h2 id="setup"><a class="header" href="#setup">Setup</a></h2> Let f(u) and f(i) denote the k dimensional embedding vector for user u and item i. Let the similarity function be s(u,i) which is typically f(u)Tf(i), and distance function d(u,i) which is typically ∣∣f(u)−f(i)∣∣22. Then some common loss functions may be denoted as below. Pointwise Losses are typically low-performing. For a given <code>(u, i)</code> pair, pointwise losses assume the presence of a <code>0, 1</code> label for relevance, and tries to predict it. The typical pointwise loss is Binary Cross Entropy, which may be expressed as: LBCE=(u,i)∈D∑logσ(s(u,i))−(u,j)∈/D∑log(1−σ(s(u,j))) Pairwise Losses assume the presence of training triplets <code>(u, i, j)</code> which correspond to user, positive item and negative item. A typical pairwise loss is Bayesian Personalized Ranking, as follows: −(u,i,j)∈τ∑logσ[s(u,i)−s(u,j)] <h2 id="weighted-matrix-factorization"><a class="header" href="#weighted-matrix-factorization">Weighted Matrix Factorization</a></h2> This describes the <a href="https://cornac.readthedocs.io/en/stable/api_ref/models.html#module-cornac.models.wmf.recom_wmf">Cornac implementation</a> of WMF. The code: <ul> <li><a href="https://github.com/PreferredAI/cornac/blob/master/cornac/models/wmf/wmf.py">wmf.py</a></li> <li><a href="https://github.com/PreferredAI/cornac/blob/master/cornac/models/wmf/recom_wmf.py">recom_wmf.py</a></li> </ul> Let A∈Rn×m describe a rating matrix of n users and m items. For simplicity, we may restrict Aij∈[0,1]. Given a user embedding matrix U∈Rn×k and item embedding matrix V∈Rm×k, WMF computes the similarity score as the dot product U⋅VT∈Rn×m. The general loss function is: L=i,j : Aij=1∑(Aij−Ui⋅VjT)2+bi,j : Aij=0∑(Aij−Ui⋅VjT)2+u⋅∣∣U∣∣F2+v⋅∣∣V∣∣F2 The idea is to simply take the squared error from the true ratings matrix as our loss, but apply a lower weightage to elements in the rating matrix where the rating is zero (as these are usually unobserved / implicit negatives that we are less confident about). Usually <code>b</code> is set to <code>0.01</code>. Regularization is performed on the user and item embedding matrices, with u∈R,v∈R as hyperparameters to adjust the strength of regularization. For cornac, this loss is adapted to the mini batch setting. Specifically, the algorithm is: <ol> <li>Draw a mini batch (default: <code>B = 128</code>) of items but use all the users</li> <li>Compute the model predictions P=U⋅VbatchT∈Rn×B</li> <li>Compute squared error E=(Abatch−P)2∈Rn×B</li> <li>Multiply matrix of weights (either <code>1</code> for positive ratings or <code>b</code> for negative ratings) element-wise with E</li> <li>loss=sum(E)+u⋅∣∣U∣∣F2+v⋅∣∣Vbatch∣∣F2</li> </ol> Note that Adam optimizer is used, and gradients are clipped between <code>[-5, 5]</code>. <h2 id="bilateral-variational-autoencoder-bivae"><a class="header" href="#bilateral-variational-autoencoder-bivae">Bilateral Variational Autoencoder (BiVAE)</a></h2> <a href="https://github.com/recommenders-team/recommenders/blob/main/examples/02_model_collaborative_filtering/cornac_bivae_deep_dive.ipynb">Recommenders BiVAE Deep Dive</a> <a href="https://dl.acm.org/doi/10.1145/3437963.3441759">BiVAE Paper</a> A working implementation of BiVAE is available on Cornac. A variational autoencoder improves over traditional linear matrix factorization methods by using non-linearity and a probabilistic formulation. Given a user, the autoencoder encodes the data representing the entity into a vector in some latent space. A decoder then takes the vector in the latent space and decodes it into something close to the original data. The difference between VAE and a regular autoencoder is that it doesn't learn a fixed vector representation, but rather a probability distribution in the latent space. This allows it to model noisy, sparse interaction data better. <h2 id="splitting"><a class="header" href="#splitting">Splitting</a></h2> <code>recommenders</code> uses a few different types of data splitting: <ol> <li>Stratified spltting.</li> </ol> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="evaluation"><a class="header" href="#evaluation">Evaluation</a></h1> Evaluation is a non-trivial topic for recsys, and different approaches measure different things. Suppose we have a dataset of user-item interactions with a timestamp. Random splitting simply takes a random split of say 75% for train and 25% for test. The problems with this approach: <ul> <li>No guarantee of overlap in users across train and test. If a user does not appear in the train set, it is not possible to recommend items for him/her in the test set.</li> <li>Chronological overlap between train and test set, leading to data leakage issues.</li> </ul> Stratified splitting addresses the user overlap issue by ensuring that the number of rows per user in the train and test set are approximately 75% and 25% of the number of rows in the original data respectively. This ensures that we have sufficient training and test data for each user, so that the collaborative filtering algorithm has a fair chance of recommending items for each user. However, stratified splitting still involves randomly assigning rows for each user into the train and test set, which does not address the chronological overlap issue. Temporal stratified splitting addresses this issue by assigning the 75% and 25% of train and test data based on chronological order. In other words, the oldest 75% of data for each user is assigned to the train set. The extreme version of temporal stratified splitting is leave last out splitting, in which all but the latest row for each user is put into the train set. This is suitable for settings where the task is to predict the very next action which the user will take (e.g. which song will the user listen to next). Note that temporal stratified splitting may potentially introduce temporal overlap between the train and test sets across users. That is, the train set period for user A may potentially overlap with the test set period for user B. Hence, if there are strong concerns with temporal effects in the dataset, we may need to be mindful of this. Global temporal splitting addresses this issue by assigning the oldest 75% of data across all users to the train set. This addresses the data leakage issue and more closely resembles actual production setting. However, there is no guarantee on the amount of train/test data for each user. Hence we may need to drop rows where there exists test data for user A but no corresponding train data due to the global temporal cutoff. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="ab-testing"><a class="header" href="#ab-testing">AB Testing</a></h1> <h2 id="references-1"><a class="header" href="#references-1">References</a></h2> <ul> <li> General: <ul> <li><a href="https://medium.com/airbnb-engineering/experiments-at-airbnb-e2db3abf39e7">Experiments at Airbnb</a></li> <li><a href="https://www.evanmiller.org/ab-testing/">Evan Miller's AB Testing Tools</a></li> <li><a href="https://posthog.com/product-engineers/ab-testing-examples">PostHog AB Testing Examples</a></li> <li><a href="https://monzo.com/blog/2022/05/24/pellets-not-cannonballs-how-we-experiment-at-monzo">How we experiment at Monzo</a></li> <li><a href="https://drive.google.com/file/d/1oK2HpKKXeQLX6gQeQpfEaCGZtNr2kR76/view">Ron Kohavi 2022 - AB Testing Intuition Busters</a></li> <li><a href="https://towardsdatascience.com/a-b-testing-a-complete-guide-to-statistical-testing-e3f1db140499">Francesco Casalegno - AB Testing A Complete Guide</a></li> </ul> </li> <li> On Bayesian AB Testing: <ul> <li><a href="https://medium.com/convoy-tech/the-power-of-bayesian-a-b-testing-f859d2219d5">Convoy - The Power of Bayesian AB Testing</a></li> <li><a href="https://www.chrisstucchio.com/pubs/VWO_SmartStats_technical_whitepaper.pdf">Chris Stucchio's White Paper</a></li> <li><a href="https://medium.com/vptech/introduction-to-bayesian-a-b-testing-in-python-df81a9b3f5fd">VeepeeTech Intro to Bayesian AB Testing</a></li> </ul> </li> </ul> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="examples"><a class="header" href="#examples">Examples</a></h1> These summaries are based on reading <a href="https://posthog.com/product-engineers/ab-testing-examples">PostHog's article on AB testing</a> and studying the ones of interest further. <h2 id="ab-testing-at-airbnb"><a class="header" href="#ab-testing-at-airbnb">AB Testing at Airbnb</a></h2> AB Testing is crucial because the outside world often has a larger effect on metrics than product changes. Factors like seasonality, economy can cause metrics to fluctuate greatly, hence a controlled experiment is necessary to control for external factors and isolate the effects of the product change. Airbnb has a complex booking flow of <code>search</code> -> <code>contact</code> -> <code>accept</code> -> <code>book</code>. While they track the AB impact of each stage, the main metric of interest is the search to book metric. One pitfall is stopping the AB Test too early. Airbnb noticed that their AB tests tend to follow a pattern of hitting significance early on but returning to neutral when the test has run its full course. This is a phenomenon known as peeking, which is to repeatedly examine an ongoing AB test. It makes it much more likely to find a significant effect when there isn't, since we are doing a statistical test each time we peek. For Airbnb, they hypothesize that this is also a phenomenon caused by the long lead time it takes from <code>search</code> -> <code>book</code>, such that early converters have a disproportionately large influence at the beginning of the experiment. The natural solution to this problem is to conduct power analysis and determine the desired sample size prior to the experiment. However, Airbnb runs multiple AB tests at the same time, hence they required an automatic way to track ongoing AB tests and report when significance has been reached. Their solution back then was to create a <code>dynamic p-value graph</code>. The idea is that on day 1 of the experiment, we would require a very low p-value to declare success. As time goes on and more samples are collected, we can gradually increase the p-value until it hits 5%. The shape of this graph is unique to their platform and they did extensive simulations to create it, so this solution is not very generalizable. Another pitfall was assuming that the system is working. After running an AB test for shifting from more words to more pictures, the initial result was neutral. However, they investigated and found that most browsers had a significant positive effect except for Internet Explorer. It turned out that the change had some breaking effect on older IE browsers. After fixing that, the overall result became positive. Hence some investigation is warranted when the AB test results are unintuitive. However, one needs to be cautious of multiple-testing when investigating breakdowns, since we are conducting multiple statistical tests. Airbnb has a strong AB testing culture - only 8% of AB tests are successful (see Ron Kohavi's <a href="https://www.linkedin.com/posts/ronnyk_interesting-discussion-in-thread-below-about-activity-7114309358942900224-X26v">LinkedIn Post</a>). <h2 id="ab-testing-at-monzo"><a class="header" href="#ab-testing-at-monzo">AB Testing at Monzo</a></h2> Monzo has a big AB testing culture - they ran 21 experiments in 6 months. Monzo has a bottom-up AB testing culture where anyone can write a proposal on Notion. Some of the best ideas come from the customer operations staff working on the frontlines. A proposal comprises the following sections: <ul> <li>What problem are you trying to solve?</li> <li>Why should we solve it?</li> <li>How should we solve it (optional)?</li> <li>What is the ideal way to solve this problem (optional)?</li> </ul> Many proposals end up becoming AB experiments. Monzo prefers launching pellets rather than cannonballs. This means that each experiment comprises small changes, is quick to build, and helps the team learn quickly. <h2 id="ab-testing-at-convoy"><a class="header" href="#ab-testing-at-convoy">AB Testing at Convoy</a></h2> Convoy argues that bayesian AB testing is more efficient than frequentist AB testing and allows them to push out product changes faster while still controlling risk. The argument against frequentist AB testing is as follows. Under traditional AB testing, we define a null hypothesis using the control group (call it A), and declare a treatment (call it B) as successful if the treatment value has a significant p-value, i.e. it falls outside of the range of reasonable values under the null. Based on power analysis and an expected effect size, we predetermine the necessary sample size to achieve sufficient power, and once this sample size is reached, we have a binary <code>success</code> or <code>failure</code> result based on the p-value. Convoy argues that this approach is safe but inefficient. This is because prior to the sample size being reached, we do not have a principled way of saying anything about the effectiveness of the treatment, even if it is performing better. Furthermore, frequentist AB testing gives us a binary result, but it does not quantify the size of the difference. Specifically, an insignificant test where <code>E(A)=10%, E(B)=11%</code> is quite different from <code>E(A)=15%, E(B)=10%</code>. For the former case, one can argue for launching B even if the p-value did not hit significance, whereas for the latter we should definitely not launch. Bayesian analysis comes in to make the above intuition concrete. Suppose we are interested in the clickthrough rate (CTR) of variant A vs B. Bayesian analysis provides a distribution of the average CTR for each variant <code>A, B</code> at any point of the AB test, based on the results that it has seen thus far. These posterior distributions reflect both the mean of the data (how far apart E(A) is from E(B)) and the variance of the data (how spread out the distributions are), allowing us to quantify how much we stand to gain if we were to pick either variant A or B at this point in time. Concretely, they define a loss function as follows. Let α and β be the unobserved true CTR for variants <code>A</code> and <code>B</code> respectively, and let the variable x denote which variant we decide to choose. Then our loss for choosing each variant can be expressed as: L(α,β,x)={max(β−α, 0)if max(α−β, 0)if x=Ax=B In other words, the loss above expresses how much we stand to lose by picking the unfortunately wrong variant based on incomplete evidence at this point in time. Of course, we do not know the true values of α and β, so we need to estimate the loss using our posterior distributions which we computed from data. We then compute the expected loss based on the posterior distributions α^∼A, β^∼B as such: E(L^)=∫α^∼A∫β^∼BL(α^,β^,x)⋅f(α^,β^)dα dβ Here, f(α^,β^) is the joint posterior distribution, which I believe we can obtain by multiplying the two independent posterior distributions α^, β^ together. We can also perform random draws from the posterior distributions to estimate this statistic. Finally, we make a decision by choosing the variant that dips below a certain loss threshold, which is usually a very small value. The appeal of the bayesian approach is two-fold: <ol> <li>It allows us to make faster decisions. Suppose an experiment is wildly successful, and it is clear within a day that variant B is better. Bayesian analysis will be able to reveal this result, whereas frequentist analysis will tell us to wait longer (since we estimated the effect size to be smaller).</li> <li>It allows us to control risk. Since we are making decisions based on minimizing risk (supposing we had picked the poorer variant), we may be sure that even if we are wrong, it will not severely degrade our product. So supposing that there is no significant engineering cost between variant A and B, we can more rapidly roll out new iterations with the assurance that on average our product will be improving.</li> </ol> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="power-analysis"><a class="header" href="#power-analysis">Power Analysis</a></h1> Reference: <a href="https://towardsdatascience.com/probing-into-minimum-sample-size-formula-derivation-and-usage-8db9a556280b">Probing into Minimum Sample Size by Mintao Wei</a> How to determine the minimum sample size required to achieve a certain significance level and power desired? The following table helps us understand how type I and type II errors come into play: <div class="table-wrapper"><table><thead><tr><th style="text-align: center"></th><th style="text-align: center">Null Hypothesis: A is True</th><th style="text-align: center">Alternate Hypothesis: B is True</th></tr></thead><tbody> <tr><td style="text-align: center">Reject A</td><td style="text-align: center">Type I Error</td><td style="text-align: center">Good statistical power</td></tr> <tr><td style="text-align: center">Accept A</td><td style="text-align: center">Good significance level</td><td style="text-align: center">Type II Error</td></tr> </tbody></table> </div> Type I Error refers to rejecting the null hypothesis when it is actually true, e.g. when we think that an AA test has significant difference. In short, it means we were too eager to deploy a poor variant. This should happen with probability α, which is the significance level which we set (typically 5%). We have a better handle on type I error because the baseline conversion rate is typically known prior to an experiment. Type II Error refers to failing to reject the null hypothesis when the alternate is actually true, i.e. we failed to get a significant effect on an improvement that is known to be better. In short, we were too conservative and failed to deploy a winning variant. In order to reason about type II error, we need to make a guess on what is the distribution of test variant <code>B</code>. Typically, this is done by assuming a minimum effect δ we wish to detect, and setting μB=μA+δ, and re-using the standard deviation from <code>A</code>. With these assumptions in place, we use power=1−β to determine the type II error that should only occur with probability β (typically 20%). Note that since δ is the minimum effect we wish to detect, if the actual effect turned out to be larger, the type II error can only be smaller than our desired amount, which is ok. Now we can derive the formula for the minimum sample size required to achieve the desired levels of type I and type II error respectively. Let us define the baseline conversion rate as p, and the minimum relative detectable effect rate as d. Consequently, the minimum detectable delta is δ=d×p. Let the desired power level be 1−β, and the desired significance level as α. Assume the scenario where we are running an AA or AB test with two variants of sample size N each. Firstly, we write down the distribution of the sample mean difference supposing we knew the true population means and standard deviations. Let E(XA)=μA,Var(XA)=σA2 and E(XB)=μB,Var(XB)=σB2. Note that XA,XB may have arbitrary distributions, e.g. they could measure proportions, revenue etc. Under the central limit theorem, the sample means will be distributed like so with NA,NB samples: XˉA∼N(μA,NAσA2), XˉB∼N(μB,NBσB2). Importantly, the difference of the sample means will have the distribution below. Note that we add the variances together because Var(B−A)=Var(B)+Var(A) for any two independent random variables A,B. XˉD=XˉB−XˉA∼N(μB−μA,NAσA2+NBσB2) Now we can start working from the desired α,β levels to the minimum sample size. We need to ensure that both objectives below are achieved with our sample size NA,NB: <ol> <li>Assuming null hypothesis to be true, ensure that type I error ≤α.</li> <li>Assuming alternate hypothesis to be true, ensure that type II error ≤1−β.</li> </ol> Let us define some notation first. <ul> <li>Let z(ϕ) denote the critical value under the standard normal distribution such that P(Z≤z(ϕ))=ϕ. This is basically the <code>scipy.stats.norm.ppf</code> function, e.g. z(0.975)=1.96.</li> <li>We also want to denote the critical value under the distribution XˉD of the sample mean difference under the null or alternate hypothesis (these are non-standard normal distributions). Let these be zXˉD∣H0(ϕ) and zXˉD∣H1(ϕ) respectively.</li> </ul> <div class="table-wrapper"><table><thead><tr><th style="text-align: center"><img src="ab_test/../images/diagram_power_analysis.png" alt="Illustration for Power Analysis Derivation" /></th></tr></thead><tbody> <tr><td style="text-align: center">Illustration for Power Analysis Derivation</td></tr> </tbody></table> </div> For objective 1, assuming the null hypothesis and using equation (1) above, we have XˉD∣H0∼N(0,NAσA2+NBσB2). Since α is a two-tailed probability and we want the critical region on the right-side, let α′=1−α/2. E.g. α=0.05 implies α′=0.975. Then: zXˉD∣H0(α′)=z(α′)⋅NAσA2+NBσB2<svg xmlns="http://www.w3.org/2000/svg" width='400em' height='3.08em' viewBox='0 0 400000 3240' preserveAspectRatio='xMinYMin slice'><path d='M473,2793 c339.3,-1799.3,509.3,-2700,510,-2702 l0 -0 c3.3,-7.3,9.3,-11,18,-11 H400000v40H1017.7 s-90.5,478,-276.2,1466c-185.7,988,-279.5,1483,-281.5,1485c-2,6,-10,9,-24,9 c-8,0,-12,-0.7,-12,-2c0,-1.3,-5.3,-32,-16,-92c-50.7,-293.3,-119.7,-693.3,-207,-1200 c0,-1.3,-5.3,8.7,-16,30c-10.7,21.3,-21.3,42.7,-32,64s-16,33,-16,33s-26,-26,-26,-26 s76,-153,76,-153s77,-151,77,-151c0.7,0.7,35.7,202,105,604c67.3,400.7,102,602.7,104, 606zM1001 80h400000v40H1017.7z'/></svg> Note that equation (2) above tells us the critical value such that we will reject the null hypothesis if the sample mean of B is greater than this value. To satisfy objective 2, we must thus ensure that the probability of rejecting the null hypothesis is at least power=1−β. In other words, we want δ−zXˉD∣H1(1−β)≥zXˉD∣H0(α′). Assuming the alternate hypothesis and again using equation (1), we have XˉD∣H1∼N(δ,NAσA2+NBσB2). So then: δ−zXˉD∣H1(1−β)δ−z(1−β)×NAσA2+NBσB2<svg xmlns="http://www.w3.org/2000/svg" width='400em' height='3.08em' viewBox='0 0 400000 3240' preserveAspectRatio='xMinYMin slice'><path d='M473,2793 c339.3,-1799.3,509.3,-2700,510,-2702 l0 -0 c3.3,-7.3,9.3,-11,18,-11 H400000v40H1017.7 s-90.5,478,-276.2,1466c-185.7,988,-279.5,1483,-281.5,1485c-2,6,-10,9,-24,9 c-8,0,-12,-0.7,-12,-2c0,-1.3,-5.3,-32,-16,-92c-50.7,-293.3,-119.7,-693.3,-207,-1200 c0,-1.3,-5.3,8.7,-16,30c-10.7,21.3,-21.3,42.7,-32,64s-16,33,-16,33s-26,-26,-26,-26 s76,-153,76,-153s77,-151,77,-151c0.7,0.7,35.7,202,105,604c67.3,400.7,102,602.7,104, 606zM1001 80h400000v40H1017.7z'/></svg>≥zXˉD∣H0(α′)≥z(α′)×NAσA2+NBσB2<svg xmlns="http://www.w3.org/2000/svg" width='400em' height='3.08em' viewBox='0 0 400000 3240' preserveAspectRatio='xMinYMin slice'><path d='M473,2793 c339.3,-1799.3,509.3,-2700,510,-2702 l0 -0 c3.3,-7.3,9.3,-11,18,-11 H400000v40H1017.7 s-90.5,478,-276.2,1466c-185.7,988,-279.5,1483,-281.5,1485c-2,6,-10,9,-24,9 c-8,0,-12,-0.7,-12,-2c0,-1.3,-5.3,-32,-16,-92c-50.7,-293.3,-119.7,-693.3,-207,-1200 c0,-1.3,-5.3,8.7,-16,30c-10.7,21.3,-21.3,42.7,-32,64s-16,33,-16,33s-26,-26,-26,-26 s76,-153,76,-153s77,-151,77,-151c0.7,0.7,35.7,202,105,604c67.3,400.7,102,602.7,104, 606zM1001 80h400000v40H1017.7z'/></svg> For the purpose of getting a minimum N, we assume N=NA=NB. Then using this and squaring both sides with some rearranging gives us: N≥δ2(σA2+σB2)(z(1−β)+z(α′))2 Which gives us the required minimum sample size equation. If we assume σA=σB, as is often assumed because we do not know the variance of the treatment, then it simplifies to the following form (as seen in <a href="https://drive.google.com/file/d/1oK2HpKKXeQLX6gQeQpfEaCGZtNr2kR76/view">Ron Kohavi's paper</a>). N≥δ22σA2⋅(z(1−β)+z(α′))2 <h2 id="bernoulli-events"><a class="header" href="#bernoulli-events">Bernoulli Events</a></h2> The equation (3) above for the minimum sample size requires us to know the standard deviation under the null and alternate hypotheses. Usually, the standard deviation under the null is computed from historical data, and it is assumed that σA=σB. However, if the event we are interested in may be represented as a bernoulli random variable (e.g. an impression is shown and user either clicks or does not click with some probability), the equation may be simplified. Specifically, the variance of a bernoulli random variable with probability p is p⋅(1−p). Thus, if XA∼Bernoulli(pA), then Var(XA)=pA⋅(1−pA), and likewise for XB. So we can use σA=pA⋅(1−pA)<svg xmlns="http://www.w3.org/2000/svg" width='400em' height='1.28em' viewBox='0 0 400000 1296' preserveAspectRatio='xMinYMin slice'><path d='M263,681c0.7,0,18,39.7,52,119 c34,79.3,68.167,158.7,102.5,238c34.3,79.3,51.8,119.3,52.5,120 c340,-704.7,510.7,-1060.3,512,-1067 l0 -0 c4.7,-7.3,11,-11,19,-11 H40000v40H1012.3 s-271.3,567,-271.3,567c-38.7,80.7,-84,175,-136,283c-52,108,-89.167,185.3,-111.5,232 c-22.3,46.7,-33.8,70.3,-34.5,71c-4.7,4.7,-12.3,7,-23,7s-12,-1,-12,-1 s-109,-253,-109,-253c-72.7,-168,-109.3,-252,-110,-252c-10.7,8,-22,16.7,-34,26 c-22,17.3,-33.3,26,-34,26s-26,-26,-26,-26s76,-59,76,-59s76,-60,76,-60z M1001 80h400000v40h-400000z'/></svg> and σB=(pA+δ)⋅(1−pA−δ)<svg xmlns="http://www.w3.org/2000/svg" width='400em' height='1.28em' viewBox='0 0 400000 1296' preserveAspectRatio='xMinYMin slice'><path d='M263,681c0.7,0,18,39.7,52,119 c34,79.3,68.167,158.7,102.5,238c34.3,79.3,51.8,119.3,52.5,120 c340,-704.7,510.7,-1060.3,512,-1067 l0 -0 c4.7,-7.3,11,-11,19,-11 H40000v40H1012.3 s-271.3,567,-271.3,567c-38.7,80.7,-84,175,-136,283c-52,108,-89.167,185.3,-111.5,232 c-22.3,46.7,-33.8,70.3,-34.5,71c-4.7,4.7,-12.3,7,-23,7s-12,-1,-12,-1 s-109,-253,-109,-253c-72.7,-168,-109.3,-252,-110,-252c-10.7,8,-22,16.7,-34,26 c-22,17.3,-33.3,26,-34,26s-26,-26,-26,-26s76,-59,76,-59s76,-60,76,-60z M1001 80h400000v40h-400000z'/></svg> and substitute these into equation (3). We will then be able to have a minimum sample size formula by just specifying α, β, baseline conversion pA and minimum relative difference d. This is the formula used by <a href="https://www.evanmiller.org/ab-testing/sample-size.html">Evan Miller's sample size calculator</a>. <h2 id="imbalanced-ab-test"><a class="header" href="#imbalanced-ab-test">Imbalanced AB Test</a></h2> Another common scenario is the case where we do not split 50-50, i.e. NA=NB. In this case, suppose we have NA=p×NB=p×n, where p>1. For example, if we have a 90-10 split, then <code>p=9</code>. Then we get: δδn≥(z(α′)+z(1−β))NAσA2+NBσB2<svg xmlns="http://www.w3.org/2000/svg" width='400em' height='3.08em' viewBox='0 0 400000 3240' preserveAspectRatio='xMinYMin slice'><path d='M473,2793 c339.3,-1799.3,509.3,-2700,510,-2702 l0 -0 c3.3,-7.3,9.3,-11,18,-11 H400000v40H1017.7 s-90.5,478,-276.2,1466c-185.7,988,-279.5,1483,-281.5,1485c-2,6,-10,9,-24,9 c-8,0,-12,-0.7,-12,-2c0,-1.3,-5.3,-32,-16,-92c-50.7,-293.3,-119.7,-693.3,-207,-1200 c0,-1.3,-5.3,8.7,-16,30c-10.7,21.3,-21.3,42.7,-32,64s-16,33,-16,33s-26,-26,-26,-26 s76,-153,76,-153s77,-151,77,-151c0.7,0.7,35.7,202,105,604c67.3,400.7,102,602.7,104, 606zM1001 80h400000v40H1017.7z'/></svg>≥(z(α′)+z(1−β))pnσA2+p⋅σB2<svg xmlns="http://www.w3.org/2000/svg" width='400em' height='3.08em' viewBox='0 0 400000 3240' preserveAspectRatio='xMinYMin slice'><path d='M473,2793 c339.3,-1799.3,509.3,-2700,510,-2702 l0 -0 c3.3,-7.3,9.3,-11,18,-11 H400000v40H1017.7 s-90.5,478,-276.2,1466c-185.7,988,-279.5,1483,-281.5,1485c-2,6,-10,9,-24,9 c-8,0,-12,-0.7,-12,-2c0,-1.3,-5.3,-32,-16,-92c-50.7,-293.3,-119.7,-693.3,-207,-1200 c0,-1.3,-5.3,8.7,-16,30c-10.7,21.3,-21.3,42.7,-32,64s-16,33,-16,33s-26,-26,-26,-26 s76,-153,76,-153s77,-151,77,-151c0.7,0.7,35.7,202,105,604c67.3,400.7,102,602.7,104, 606zM1001 80h400000v40H1017.7z'/></svg>≥δ2(z(α′)+z(1−β))(pσA2+σB2) Note that n gives the sample size required for NB, and the total sample size required across both groups is n(p+1). <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="large-language-models"><a class="header" href="#large-language-models">Large Language Models</a></h1> LLMs are generally used in an auto-regressive way, where the user supplies a prompt and the LLM returns a generated response. This framework makes it amenable to a wide range of tasks. HuggingFace has a great <a href="https://huggingface.co/blog/optimize-llm">blog post</a> explaining how we can run LLMs on humble hardware. Typically, LLMs have billions of parameters. The following rules of thumb helps us know how much memory we need to load the LLM into memory. For a model with <code>X billion</code> parameters: <ul> <li>Loading in <code>float32</code> requires <code>4X GB</code> of VRAM</li> <li>Loading in <code>bfloat16</code> requires <code>2X GB</code> of VRAM</li> <li>Loading in <code>int8</code> requires <code>X GB</code> of VRAM</li> </ul> Hence we see that we can load a <code>~7 billion</code> parameters model with around <code>14GB</code> of VRAM if loaded in <code>bfloat16</code>, which makes it feasible to run on GPUs like Tesla T4 with <code>16GB</code> of VRAM. This can be done when loading the model with <code>from_pretrained(..., torch_dtype=torch.bfloat16)</code>. Most models are trained in <code>bfloat16</code> anyway, so it makes sense to load them at that precision. Current popular open source LLM of that size includes <code>mosaicml/mpt-7b</code>, which can be easily downloaded and used using huggingface. <h2 id="quantization"><a class="header" href="#quantization">Quantization</a></h2> It turns out that we can lower the precision of models even further than 16 bits if we use a quantization method (see e.g. <a href="https://arxiv.org/abs/2208.07339">Dettmers 2022</a>, this paper is the basis for the package <code>bitsandbytes</code> used for quantization). The general idea is akin to encoding - we encode each number from a higher precision into a "codeword" in the lower precision (i.e. quantization). Numbers that are close to one another in the higher precision may get mapped to the same "codeword". When we want to use the encoded value, we look up the value in the higher precision that it maps to (i.e. de-quantization). When applying this to quantizing a neural network, the steps involved are: <ol> <li>Quantize all model weights to target precision (e.g. <code>int8</code>)</li> <li>Pass the input vector at <code>bfloat16</code></li> <li>At each layer, dequantize the weights and perform matmul in <code>bfloat16</code></li> <li>Quantize the weights again for storage</li> </ol> Hence while quantization lowers the memory footprint of the model, it may increase inference time. To use quantization, we need to do the following (also make sure <code>bitsandbytes</code> is pip installed). We can also pass <code>load_in_4bit=True</code> for 4bit quantization. More info on quantization usage is available at <a href="https://huggingface.co/docs/transformers/main_classes/quantization#general-usage">HuggingFace</a>. An important note is that a GPU is required for quantization, at least in the <code>bitsandbytes</code> package. <pre><code class="language-python">model = AutoModelForCausalLM.from_pretrained(..., load_in_8bit=True) </code></pre> <h2 id="flash-attention"><a class="header" href="#flash-attention">Flash Attention</a></h2> The self-attention mechanism (<a href="https://arxiv.org/abs/2205.14135">Dao 2022</a>) is at the heart of the transformer performance but is also a major bottleneck in terms of memory and computational cost. One of the most successful optimizations for the attention mechanism is the Flash Attention paper. Suppose we have an input sequence of embeddings X=(x1,...,xN) where xi∈Rk, such that X∈Rk×N. The transformer stores parameters Wq,Wk,Wv∈Rk×d, such that Q=XT⋅Wq,K=XT⋅Wk,V=XT⋅Wv such that Q,K,V∈RN×d. The self-attention matrix S=QKT∈RN×N is then computed to represent the pairwise interaction between tokens at position i (ith row) and position j (jth column). The row-wise softmax is taken P=softmax(S) to convert these into probabilities and finally the output is O=P⋅V. Typically, N is much larger than the hidden dimensions k,d, as N can be <code>2,048</code> or larger. Hence the QKT∈RN×N matrix is the bottleneck for memory and computation. The flash attention proposes to do this computation in a block-wise manner to reduce the memory usage. Furthermore, the algorithm also speeds up the computation compared to naive attention because the block-wise implementation minimizes the number of read-write operations between the faster SRAM and slower HBM. More details can be found in the notebook at <a href="llm/./papers/dao_2022.html">Dao 2022 - FlashAttention</a>. We can utilize flash attention like so: <pre><code class="language-python">%pip install optimum model.to_bettertransformer() </code></pre> Note that this is only supported for models that have implemented flash attention, e.g. <code>gpt-neox</code>, <code>bloom</code> etc. Flash Attention is now support natively within Pytorch as <code>torch.nn.functional.scaled_dot_product_attention</code> (see <a href="https://pytorch.org/blog/out-of-the-box-acceleration/">blog</a>). The usage is like below. We need <code>transformers>=4.36</code> and <code>torch>=2.1.1</code> to use it. <pre><code class="language-python">from optimum.bettertransformer import BetterTransformer from transformers import AutoModelForCausalLM model = AutoModelForCausalLM.from_pretrained("gpt2-large", torch_dtype=torch.float16) model = BetterTransformer.transform(model, keep_original_model=False) </code></pre> <h2 id="position-representation"><a class="header" href="#position-representation">Position Representation</a></h2> Recent innovations in position encoding has led to accuracy improvements for long input sequences. The initial attention papers used absolute position embeddings. Given an input sequence of embeddings X=(x1,...,xN)∈Rd×N, absolute position embeddings p1,...,pN∈Rd are generated by the model. These position embeddings are added to the input sequence X^=(xi+p1,...,xN+pN), thereby allowing to model to use these position cues. It turns out that fixed positional embeddings are not ideal because they require the model to learn a fixed, unique representation of each position 1,...,N. This does not represent language well, because the word in position i in one sentence does not necessarily serve the same purpose as a word in the same position in another sentence. Rather, it is the relative distance between words that we want to encode in the model. Furthermore, training absolute position embeddings makes it difficult for our model to generalize to texts with longer sequences than what it was trained with. Recent papers advocate for relative positional embeddings, with the following differences: <ul> <li>Relative position encoding rather than absolute position encoding</li> <li>The encoding of relative position is done most naturally within the QKT self-attention matrix, since that is where the relative degree of interaction between tokens at different positions is encoded</li> <li>The encoding should be such that tokens further apart have a lower value in the self-attention matrix and tokens closer together have a higher value</li> </ul> Rotational Position Embeddings (<code>RoPE</code>) (<a href="https://arxiv.org/pdf/2104.09864.pdf">Su 2021</a>) proposes rotating the query and key vectors by an angle proportional to the relative distance between the two positions. Specifically q^iTk^j=qiTRθ,i−jkj, where Rθ,i−j is a rotational matrix that performs the rotation. Attention with Linear Biases (<code>ALiBi</code>) (<a href="https://arxiv.org/pdf/2108.12409.pdf">Press 2022</a>) proposes an even simpler method. It simply subtracts ∣i−j∣/m from row i and column j of the self-attention matrix, where m is a fixed scalar (specific to each attention head). Intuitively, it penalizes the attention proportional to the distance between the tokens. The study shows that this method outperforms RoPE as we extrapolate to longer sequences, and is conceptually simpler. <h2 id="key-value-cache"><a class="header" href="#key-value-cache">Key-Value Cache</a></h2> Most LLMs work in an auto-regressive manner, i.e. we provide an input sequence, generate the next token with the LLM, then append this token to the input sequence for the next iteration. Most LLMs are also trained with the causal language modelling objective and mask the upper triangle of the self-attention matrix, so that each query token qi can only interact with key token kj and value token vj if j≥i. This setup encourages us to cache results from previous time steps, since a lot of computation is repeated. The following is based on how I imagine this to work, after reading <a href="https://www.linkedin.com/posts/cameron-r-wolfe-ph-d-04744a238_friday-ai-fundamentals-the-kv-cache-using-activity-7095825756928311296-7xXD">Cameron R. Wolfe's LinkedIn post</a>. During training, we compute the projections Q,K,V∈RN×d, where N is the maximum sequence length and d is the hidden dimension. The final output O=softmax(QKT)V∈RN×d actually provides a d-dimension representation of the model's prediction at each of the N positions. For next token generation, we can add a projection head, say Wp∈Rd×p, where p represents the size of the vocabulary, such that A=OWp∈RN×p can represent the activations at each of the N positions for the next token. Specifically, A[0, :] represents the predictions of position 1 given input tokens [0], A[1, :] represents the predictions of position 2 given input tokens [0,1], and so on. These activations will then be fed into some cross-entropy loss such that activations at the correct token for each position gets rewarded. This allows us to do efficient training, since we simultaneously provide losses for the prediction at each of the N positions to the model for backpropagation. However, when we are doing inference generation, we only need to predict for the final position of the input sequence (suppose it is position c), i.e. we are only interested in A[c, :] and O[c, :]. Hence for starters, we only need qc:=Q[c, :] instead of the entire Q matrix, since only that row comes into play. However, we still need the entire K and V matrices, since we want qc to interact with all tokens in the input sequence. This is where the KV cache comes in - we cache the existing K and V matrices, so that we only need to project the final token of the input sequence xcTWk∈R1×d and xcTWv∈R1×d at each step and append it to the existing cached K and V. We can then compute Oc=softmax(qcKT)V∈R1×d. As one can imagine, this saves a lot of computation, but also increases memory costs. <a href="https://arxiv.org/pdf/2309.06180.pdf">Kwon 2023 - PagedAttention</a> shows that serving a <code>13B</code> model on NVIDIA A100 with 40GB of memory: <ul> <li>65% of memory is model parameters</li> <li>>30% is the KV cache</li> <li>A small amount of memory is used ephemerally for activation</li> </ul> The usage of KV cache is like so: <pre><code class="language-python">model = AutoModelForCausalLM.from_pretrained(...) model.generate(..., use_cache=True) </code></pre> <h2 id="how-to-fine-tune-an-llm"><a class="header" href="#how-to-fine-tune-an-llm">How to fine-tune an LLM</a></h2> <ul> <li><a href="https://huggingface.co/blog/trl-peft">trl RL example</a> <ul> <li>Fine tune a 20B GPT model on text generation on IMDB dataset (loaded in 8 bit)</li> <li>Since step 1 used PEFT, we need to merge the adapter weights with the base model</li> <li>Finally, use RLHF to generate positive movie reviews. They used a BERT IMDB sentiment classifer to generate rewards</li> </ul> </li> <li><a href="https://www.datacamp.com/tutorial/fine-tuning-llama-2">DataCamp example</a> <ul> <li>Using <code>SFTTrainer</code> from the <code>trl</code> library to do supervised fine-tuning</li> </ul> </li> <li><a href="https://github.com/huggingface/peft">PEFT - based on LORA</a> - PEFT is built by hugginface to support LORA.</li> </ul> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="kv-cache"><a class="header" href="#kv-cache">KV-Cache</a></h1> This article walks through a simple but rigorous demonstration of KV caching in transformer autoregressive decoding. <h2 id="the-attention-mechanism"><a class="header" href="#the-attention-mechanism">The Attention Mechanism</a></h2> In autoregressive generation, we generate tokens one at a time. Given a sequence of input embeddings [x1,x2,...,xt], we feed the embedding matrix of size B×L×d into the model to predict xt+1. Note that: <ul> <li>B denotes batch size</li> <li>L denotes sequence length</li> <li>d denotes model dimension</li> <li>V denotes output vocabulary size</li> </ul> We will focus on the self-attention block (where the KV-cache is pertinent). For exposition, we will omit the batch size dimension for now and consider a single batch. The self attention block computes: Attention(Q,K,V)=softmax(dk<svg xmlns="http://www.w3.org/2000/svg" width='400em' height='1.08em' viewBox='0 0 400000 1080' preserveAspectRatio='xMinYMin slice'><path d='M95,702 c-2.7,0,-7.17,-2.7,-13.5,-8c-5.8,-5.3,-9.5,-10,-9.5,-14 c0,-2,0.3,-3.3,1,-4c1.3,-2.7,23.83,-20.7,67.5,-54 c44.2,-33.3,65.8,-50.3,66.5,-51c1.3,-1.3,3,-2,5,-2c4.7,0,8.7,3.3,12,10 s173,378,173,378c0.7,0,35.3,-71,104,-213c68.7,-142,137.5,-285,206.5,-429 c69,-144,104.5,-217.7,106.5,-221 l0 -0 c5.3,-9.3,12,-14,20,-14 H400000v40H845.2724 s-225.272,467,-225.272,467s-235,486,-235,486c-2.7,4.7,-9,7,-19,7 c-6,0,-10,-1,-12,-3s-194,-422,-194,-422s-65,47,-65,47z M834 80h400000v40h-400000z'/></svg>QKT)V Where: <ul> <li>Q=X⋅WQ (queries)</li> <li>K=X⋅WK (keys)</li> <li>V=X⋅WV (values)</li> </ul> The transformer produces an output of shape L×d at the last layer. The output at the final position (call it ht[L]∈Rd) is finally used to generate the logits for each token to predict like so: logits=ht[L]⋅Wout∈RV <h3 id="worked-example"><a class="header" href="#worked-example">Worked Example</a></h3> Let's make this concrete. Suppose we have: <ul> <li>2 tokens already: x1,x2 (each is a 3-dimensional embedding vector)</li> <li>Projection dimension dk=2</li> </ul> Let's setup our input and projection matrices. Notice that each row of X corresponds to one position in the sequence. X=[—x1——x2—]∈R2×3,WQ,WK,WV∈R3×2 Step 1: Generate token 3 (predict x3 given [x1,x2]) We compute Q, K, V for all positions: Q=X⋅WQ=[—q1——q2—],K=X⋅WK=[—k1——k2—],V=X⋅WV=[—v1——v2—] Let's write out the full attention computation (omitting softmax and scaling for clarity): A~=Q⋅KT=[q1⋅k1q2⋅k1q1⋅k2q2⋅k2] After applying the causal mask (setting future positions to −∞) and softmax, we get attention weights A: A=[a11a210a22] Importantly, observe that the bottom row of A only depends on q2,k1 and k2. q1 is not used at all. We will see that q1 is in fact not needed to generate the next token shortly. Finally, the output is O=A⋅V: O=[a11a210a22][v11v21v12v22]=[a11v11a21v11+a22v21a11v12a21v12+a22v22] To predict x3, we only need the bottom row of O. But this accordingly only uses the bottom row of A, which only depends on q2. Hence, we can simplify the computation of the bottom row of O as follows: o2=softmax(q2⋅KT)⋅V Summary: To decode for step t, we need: <ul> <li>Only q2 (the query for the last position)</li> <li>All of K1:2</li> <li>All of V1:2</li> </ul> Now suppose we want to generate token 4. We will need: <ul> <li>Q3</li> <li>K1:3</li> <li>V1:3</li> </ul> <blockquote> Key Insight: Due to the causal mask, K1:2 and V1:2 did not change. Hence we can cache them and just append K3 and V3 to get K1:3,V1:3. This is the main idea of the KV Cache. </blockquote> <h2 id="naive-implementation-no-cache"><a class="header" href="#naive-implementation-no-cache">Naive Implementation (No Cache)</a></h2> Without caching, at each generation step we recompute K and V for the entire sequence: <pre><code class="language-python">import torch import torch.nn as nn import torch.nn.functional as F def naive_attention(X, W_q, W_k, W_v): """ X: (batch, seq_len, d_model) - full sequence Returns: (batch, seq_len, d_model) """ Q = X @ W_q # (batch, seq_len, d_k) K = X @ W_k # (batch, seq_len, d_k) V = X @ W_v # (batch, seq_len, d_v) d_k = K.shape[-1] scores = (Q @ K.transpose(-2, -1)) / (d_k ** 0.5) # Causal mask: prevent attending to future tokens seq_len = X.shape[1] mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool() scores = scores.masked_fill(mask, float('-inf')) attn = F.softmax(scores, dim=-1) return attn @ V def naive_generate(prompt_embeddings, W_q, W_k, W_v, W_out, num_tokens): """ Generate tokens without KV cache. Each step recomputes attention over the ENTIRE sequence. """ X = prompt_embeddings # (1, prompt_len, d_model) for step in range(num_tokens): # Recompute attention for ALL positions (wasteful!) attn_out = naive_attention(X, W_q, W_k, W_v) # Get logits for last position only last_hidden = attn_out[:, -1, :] # (1, d_model) logits = last_hidden @ W_out # (1, vocab_size) # Sample next token (simplified: just argmax) next_token_id = logits.argmax(dim=-1) # Embed next token and append to sequence # (In practice, look up embedding; here we simulate) next_embedding = torch.randn(1, 1, X.shape[-1]) X = torch.cat([X, next_embedding], dim=1) print(f"Step {step+1}: Computed K,V for {X.shape[1]} positions") return X </code></pre> Complexity without cache: For generating T new tokens with prompt length P: <ul> <li>Step 1: Compute K,V for P+1 positions</li> <li>Step 2: Compute K,V for P+2 positions</li> <li>...</li> <li>Step T: Compute K,V for P+T positions</li> </ul> Total K,V computations: ∑i=1T(P+i)=T⋅P+2T(T+1)=O(T2+TP) <h2 id="with-kv-cache"><a class="header" href="#with-kv-cache">With KV Cache</a></h2> The KV cache stores previously computed keys and values. At each step, we only compute K,V for the new token and concatenate with the cache: <pre><code class="language-python">def cached_attention(x_new, W_q, W_k, W_v, kv_cache): """ x_new: (batch, 1, d_model) - ONLY the new token kv_cache: tuple of (K_cache, V_cache) or None K_cache: (batch, prev_seq_len, d_k) V_cache: (batch, prev_seq_len, d_v) Returns: attention output, updated cache """ # Compute Q, K, V for the NEW token only q_new = x_new @ W_q # (batch, 1, d_k) k_new = x_new @ W_k # (batch, 1, d_k) v_new = x_new @ W_v # (batch, 1, d_v) if kv_cache is None: K = k_new V = v_new else: K_cache, V_cache = kv_cache # Append new K,V to cache K = torch.cat([K_cache, k_new], dim=1) # (batch, seq_len, d_k) V = torch.cat([V_cache, v_new], dim=1) # (batch, seq_len, d_v) # Attention: new query attends to ALL keys d_k = K.shape[-1] scores = (q_new @ K.transpose(-2, -1)) / (d_k ** 0.5) # (batch, 1, seq_len) # No causal mask needed: q_new is at position seq_len, # and we only have K,V up to seq_len (no future tokens) attn = F.softmax(scores, dim=-1) out = attn @ V # (batch, 1, d_v) # Return output and updated cache return out, (K, V) def cached_generate(prompt_embeddings, W_q, W_k, W_v, W_out, num_tokens): """ Generate tokens WITH KV cache. Two phases: 1. Prefill: Process entire prompt at once, build initial cache 2. Decode: Generate one token at a time, updating cache incrementally """ # === PREFILL PHASE === # Process entire prompt to build initial KV cache X = prompt_embeddings # (1, prompt_len, d_model) K_cache = X @ W_k # (1, prompt_len, d_k) V_cache = X @ W_v # (1, prompt_len, d_v) # Compute attention for prompt (need full Q for all positions) Q = X @ W_q d_k = K_cache.shape[-1] scores = (Q @ K_cache.transpose(-2, -1)) / (d_k ** 0.5) seq_len = X.shape[1] mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool() scores = scores.masked_fill(mask, float('-inf')) attn = F.softmax(scores, dim=-1) attn_out = attn @ V_cache last_hidden = attn_out[:, -1:, :] # Keep dim for next iteration kv_cache = (K_cache, V_cache) print(f"Prefill: Built cache with {K_cache.shape[1]} positions") # === DECODE PHASE === for step in range(num_tokens): # Get logits from last hidden state logits = last_hidden.squeeze(1) @ W_out next_token_id = logits.argmax(dim=-1) # Embed next token (simulated) next_embedding = torch.randn(1, 1, X.shape[-1]) # Compute attention with cache - only process NEW token! last_hidden, kv_cache = cached_attention( next_embedding, W_q, W_k, W_v, kv_cache ) print(f"Step {step+1}: Computed K,V for 1 position (cache size: {kv_cache[0].shape[1]})") return kv_cache </code></pre> Complexity with cache: For generating T new tokens with prompt length P: <ul> <li>Prefill: Compute K,V for P positions (once)</li> <li>Each decode step: Compute K,V for 1 position</li> </ul> Total K,V computations: P+T=O(P+T) <h2 id="comparison"><a class="header" href="#comparison">Comparison</a></h2> <div class="table-wrapper"><table><thead><tr><th>Aspect</th><th>Without Cache</th><th>With Cache</th></tr></thead><tbody> <tr><td>K,V computations per step</td><td>O(seq_len)</td><td>O(1)</td></tr> <tr><td>Total K,V computations</td><td>O(T2+TP)</td><td>O(P+T)</td></tr> <tr><td>Memory</td><td>O(d)</td><td>O((P+T)⋅d)</td></tr> <tr><td>Attention computation</td><td>O(seq_len2)</td><td>O(seq_len) per step</td></tr> </tbody></table> </div> The KV cache trades memory for computation. This is almost always worthwhile because: <ol> <li>Matrix multiplications (K,V projections) are expensive</li> <li>Memory is typically available</li> <li>The speedup is dramatic for long sequences</li> </ol> <h2 id="complete-working-example"><a class="header" href="#complete-working-example">Complete Working Example</a></h2> <pre><code class="language-python">import torch import torch.nn.functional as F torch.manual_seed(42) # Dimensions batch_size = 1 d_model = 64 d_k = d_v = 32 vocab_size = 100 prompt_len = 5 num_generate = 3 # Initialize weights W_q = torch.randn(d_model, d_k) / (d_model ** 0.5) W_k = torch.randn(d_model, d_k) / (d_model ** 0.5) W_v = torch.randn(d_model, d_v) / (d_model ** 0.5) W_out = torch.randn(d_v, vocab_size) / (d_v ** 0.5) # Simulate prompt embeddings prompt = torch.randn(batch_size, prompt_len, d_model) print("=== WITHOUT KV CACHE ===") naive_generate(prompt.clone(), W_q, W_k, W_v, W_out, num_generate) print("\n=== WITH KV CACHE ===") cached_generate(prompt.clone(), W_q, W_k, W_v, W_out, num_generate) </code></pre> Output: <pre><code>=== WITHOUT KV CACHE === Step 1: Computed K,V for 6 positions Step 2: Computed K,V for 7 positions Step 3: Computed K,V for 8 positions === WITH KV CACHE === Prefill: Built cache with 5 positions Step 1: Computed K,V for 1 position (cache size: 6) Step 2: Computed K,V for 1 position (cache size: 7) Step 3: Computed K,V for 1 position (cache size: 8) </code></pre> <h2 id="why-only-cache-k-and-v-not-q"><a class="header" href="#why-only-cache-k-and-v-not-q">Why Only Cache K and V (Not Q)?</a></h2> During autoregressive generation: <ul> <li>Query (Q): We only need the query for the current position to compute attention scores. Past queries are never reused.</li> <li>Keys (K): The current query must attend to all past keys. These are reused every step.</li> <li>Values (V): Once we have attention weights, we aggregate all past values. These are reused every step.</li> </ul> Mathematically, for generating token at position t: outt=softmax(dk<svg xmlns="http://www.w3.org/2000/svg" width='400em' height='1.08em' viewBox='0 0 400000 1080' preserveAspectRatio='xMinYMin slice'><path d='M95,702 c-2.7,0,-7.17,-2.7,-13.5,-8c-5.8,-5.3,-9.5,-10,-9.5,-14 c0,-2,0.3,-3.3,1,-4c1.3,-2.7,23.83,-20.7,67.5,-54 c44.2,-33.3,65.8,-50.3,66.5,-51c1.3,-1.3,3,-2,5,-2c4.7,0,8.7,3.3,12,10 s173,378,173,378c0.7,0,35.3,-71,104,-213c68.7,-142,137.5,-285,206.5,-429 c69,-144,104.5,-217.7,106.5,-221 l0 -0 c5.3,-9.3,12,-14,20,-14 H400000v40H845.2724 s-225.272,467,-225.272,467s-235,486,-235,486c-2.7,4.7,-9,7,-19,7 c-6,0,-10,-1,-12,-3s-194,-422,-194,-422s-65,47,-65,47z M834 80h400000v40h-400000z'/></svg>qt⋅[k1,...,kt]T)⋅[v1,...,vt] Only qt is new. All ki,vi for i<t were computed in previous steps. <h2 id="multiple-layers-each-layer-has-its-own-cache"><a class="header" href="#multiple-layers-each-layer-has-its-own-cache">Multiple Layers: Each Layer Has Its Own Cache</a></h2> Real transformers stack L self-attention layers. Each layer maintains its own KV cache because: <ol> <li>Each layer has different projection weights: Layer ℓ has its own WK(ℓ) and WV(ℓ)</li> <li>Each layer receives different input: The input to layer ℓ is the output of layer ℓ−1</li> </ol> Therefore, even for the same token position, the K and V values differ across layers. <pre><code class="language-python">def multi_layer_cached_generate(prompt_embeddings, layers, W_out, num_tokens): """ layers: list of dicts, each with W_q, W_k, W_v for that layer KV cache structure: list of (K_cache, V_cache) tuples, one per layer """ num_layers = len(layers) X = prompt_embeddings # (1, prompt_len, d_model) # === PREFILL: Build KV cache for ALL layers === kv_caches = [] hidden = X for layer_idx, layer in enumerate(layers): W_q, W_k, W_v = layer['W_q'], layer['W_k'], layer['W_v'] # Compute K, V for this layer's input K = hidden @ W_k # (1, prompt_len, d_k) V = hidden @ W_v # (1, prompt_len, d_v) kv_caches.append((K, V)) # Compute attention output (with causal mask) Q = hidden @ W_q d_k = K.shape[-1] scores = (Q @ K.transpose(-2, -1)) / (d_k ** 0.5) seq_len = hidden.shape[1] mask = torch.triu(torch.ones(seq_len, seq_len), diagonal=1).bool() scores = scores.masked_fill(mask, float('-inf')) attn = F.softmax(scores, dim=-1) hidden = attn @ V # Output becomes input to next layer last_hidden = hidden[:, -1:, :] print(f"Prefill: Built {num_layers} KV caches, each with {X.shape[1]} positions") # === DECODE: Update each layer's cache incrementally === for step in range(num_tokens): logits = last_hidden.squeeze(1) @ W_out next_token_id = logits.argmax(dim=-1) next_embedding = torch.randn(1, 1, X.shape[-1]) # Forward through all layers, updating each cache hidden = next_embedding new_caches = [] for layer_idx, layer in enumerate(layers): W_q, W_k, W_v = layer['W_q'], layer['W_k'], layer['W_v'] K_cache, V_cache = kv_caches[layer_idx] # Compute K, V for NEW token only at this layer q_new = hidden @ W_q # (1, 1, d_k) k_new = hidden @ W_k # (1, 1, d_k) v_new = hidden @ W_v # (1, 1, d_v) # Append to this layer's cache K = torch.cat([K_cache, k_new], dim=1) V = torch.cat([V_cache, v_new], dim=1) new_caches.append((K, V)) # Attention: new query attends to all cached K, V d_k = K.shape[-1] scores = (q_new @ K.transpose(-2, -1)) / (d_k ** 0.5) attn = F.softmax(scores, dim=-1) hidden = attn @ V # (1, 1, d_v) - input to next layer kv_caches = new_caches last_hidden = hidden cache_size = kv_caches[0][0].shape[1] print(f"Step {step+1}: Updated {num_layers} caches (each now size {cache_size})") return kv_caches </code></pre> Memory for KV cache with L layers: KV cache size=2×L×(seq_len)×dk×batch_size The factor of 2 is for K and V. For large models (many layers, large dk), this becomes the dominant memory cost during inference. Example: Llama 2 70B has: <ul> <li>L=80 layers</li> <li>dk=128 per head, 64 heads → 8192 total</li> <li>For a 4096-token sequence in fp16: 2×80×4096×8192×2 bytes≈10.7 GB</li> </ul> This is why KV cache compression techniques (quantization, eviction, etc.) are active research areas. <h2 id="references-2"><a class="header" href="#references-2">References</a></h2> <ul> <li><a href="https://huggingface.co/blog/not-lain/kv-caching">Huggingface Blog: KV Caching</a></li> </ul> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="automatic-prompt-optimization"><a class="header" href="#automatic-prompt-optimization">Automatic Prompt Optimization</a></h1> Automatic prompt optimization is a research area of recent interest. This trend arises from the following observations: <ul> <li>Adjusting textual prompts can very significantly change LLM's ability to accurately complete tasks</li> <li>The correlation between prompt and performance is not always the most obvious or intuitive to humans</li> <li>The best prompt varies depending on the LLM model (or even different iterations of the same model)</li> </ul> LLMs are increasingly used to solve diverse problems by varying the instructions. This is often a simpler solution than maintaining many models that do various specific tasks. For example, in the recsys world, companies like <a href="https://arxiv.org/abs/2501.16450">LinkedIn</a> and <a href="https://netflixtechblog.com/foundation-model-for-personalized-recommendation-1a0bd8e02d39">Netflix</a> are moving toward foundational LLM models that can replace a collection of traditional recommender models. By adjusting the instructions to the foundation model, it can achieve comparable or even better performance to these task-specific models due to favourable scaling laws with the size of the model. However, manually tuning the instruction for each task is brittle and model performance can easily change as the underlying foundation model is continually trained and improved, or swapped out. Also, we may often want to adapt a foundational model to a new task without performing further fine-tuning or training adapter weights. Hence, automatic prompt optimization becomes an important tool to maintain model performance in an operational setting. This note focuses on gradient-free methods to optimize black-box LLMs, as they are often simpler and are applicable to LLM usage via API, which is still the most common LLM usage method outside of big tech companies. I also follow <a href="llm/prompt_optimization.html#wolfe2024">Wolfe 2024</a> heavily but simplify quite a bit for readability. <h2 id="main-idea"><a class="header" href="#main-idea">Main Idea</a></h2> The main idea of automatic prompt optimization is to simply treat the task of prompt optimization as a standard machine learning problem. That is, we construct a training dataset and optimize our prompt to improve performance on it, and validate the effectiveness of the prompt on a held out validation set. The only difference is that we have to be creative in how we optimize the prompt, since we cannot use gradient descent to do so. Specifically, the general setup that we consider is such: <ul> <li>We have a dataset of inputs and labels that can be split into train and validation sets <ul> <li>Labels are optional if we are using another LLM as judge</li> </ul> </li> <li>For each training example, we can produce textual output from an LLM call: output=LLM(prompt,input)</li> <li>We have some evaluation function that returns a score for this example instance: score=Eval(output,label) <ul> <li>The evaluation function obviously depends on the specific task</li> <li>E.g. for simple tasks, the evaluation function can simply be the accuracy of output==label</li> <li>The evaluation function can also be another LLM for ambiguous tasks</li> </ul> </li> <li>We can thus compute the mean score across a set of training instances to evaluate the effectiveness of a given prompt</li> <li>At each iteration, we try one or more new prompt(s) and evaluate their performance, then have some way of generating new prompts to try again</li> <li>At the end, the best prompt is selected</li> </ul> <h2 id="methods"><a class="header" href="#methods">Methods</a></h2> Here we dive into the different methods. Some are quite ingenious in how they use LLMs to elicit the best instructions. Note that we omit papers that perform prompt optimization through some gradient-based optimization (e.g. using reinforcement learning), as these cannot be used with LLMs accessed through API. <h3 id="instruction-induction"><a class="header" href="#instruction-induction">Instruction Induction</a></h3> <a href="llm/prompt_optimization.html#honovich2022">Honovich 2022</a> is an early paper that suggests taking a random subset of training <code>(input, label)</code> pairs, showing them to an LLM, and asking the LLM to guess the instruction that produces the label from the inputs. For example, the LLM prompt may look something like: <pre><code class="language-bash">Here are the input-output pairs: Input: As soon as you can. Output: At your earliest convenience. ... The instruction was <PLS FILL IN> </code></pre> The LLM may then guess something like <code>The instruction was translate the inputs into more formal language</code>. We can then use this instruction as the new optimized prompt. Note that this method is a one step process, i.e. it does not iterate for further improvements. But I suppose we can sample multiple random subsets of training instances and generate prompts for each before picking the best prompt. <h2 id="bibliography"><a class="header" href="#bibliography">Bibliography</a></h2> <ul> <li><a id="honovich2022"></a><a href="https://arxiv.org/abs/2205.10782">Honovich 2022: Instruction Induction</a></li> <li><a id="yang2023"></a><a href="https://arxiv.org/abs/2309.03409">Yang 2023: Large Language Models as Optimizers (OPRO)</a></li> <li><a id="wolfe2024"></a><a href="https://cameronrwolfe.substack.com/p/automatic-prompt-optimization">Wolfe 2024: Automatic Prompt Optimization</a></li> </ul> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="fine-tuning"><a class="header" href="#fine-tuning">Fine-tuning</a></h1> LLMs are typically trained with next-token prediction task on large amounts of text in an unsupervised manner. Some of the behaviour in these texts are not desirable to imitate. For example, while Github is full of code repositories with common programming mistakes, we do not want the LLM to replicate such behaviour. Hence, a process of alignment is necessary to encourage the model to produce desired responses. There are typically two stages to this fine-tuning: Supervised Fine-Tuning and Reinforcement Learning from Human Feedback. Supervised Fine-Tuning (SFT). In this stage, pairs of <code>(prompt, desired response)</code> are provided to the LLM. The desired responses are often called "demonstrations" by human annotators. Some form of cross-entropy loss is then used to update the LLM to encourage it to generate the desired response. This is a straightforward approach that is similar to the next-token prediction task (except we are predicting the desired response given the prompt). In the <a href="https://arxiv.org/pdf/2203.02155.pdf">InstructGPT</a> paper, the authors show that an SFT-aligned <code>1.3B</code> model (using ~<code>13k</code> training data) generates human-preferred outputs compared to a <code>175B</code> GPT model, showing the importance of SFT. The problem with SFT is that it trains the model to provide a very specific response to a particular prompt, which makes it hard for the model to generalize. It also fails to express the natural ambiguity of responses for an open-ended prompt (e.g. <code>write a poem about AI</code>). Hence, unless we can collect quality responses for a very wide variety of prompts, SFT is limited in its ability to generalize. Reinforcement Learning from Human Feedback (RLHF). This is where RLHF comes in: given triplets of <code>(prompt, preferred response, not preferred response)</code>, we train the model to generate the preferred response in a more generalizable way. More recently, a method called Direct Preference Optimization is used to solve this problem without needing to do reinforcement learning. More on that <a href="llm/../papers/rafailov_2023.html">here</a>. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="useful-models"><a class="header" href="#useful-models">Useful Models</a></h1> There are too many models on Huggingface, so I try to keep track of useful ones here. Generally the <a href="https://huggingface.co/spaces/HuggingFaceH4/open_llm_leaderboard">open_llm_leaderboard</a> is a good place to start. <div class="table-wrapper"><table><thead><tr><th style="text-align: center">Model Name</th><th style="text-align: center"># Params (Billions)</th><th style="text-align: center">openllm score</th><th style="text-align: center">Comments</th></tr></thead><tbody> <tr><td style="text-align: center"><a href="https://huggingface.co/upstage/SOLAR-10.7B-Instruct-v1.0">SOLAR-10.7B-Instruct-v1.0</a></td><td style="text-align: center">10.7B</td><td style="text-align: center">74.2</td><td style="text-align: center">This model was trained using Supervised Finetuning (SFT) and Direct Preference Optimization (DPO) for single-turn chat. Seems to be derived from Llama. Verified that the outputs are sensible.</td></tr> </tbody></table> </div> <h2 id="testing-solar-107b"><a class="header" href="#testing-solar-107b">Testing SOLAR-10.7B</a></h2> Here are some results from testing out <code>SOLAR-10.7B-Instruct-v1.0</code> on Google Colab, with a T4 GPU. First we install packages and load the model. Note that we are quantizing it to <code>4bits</code> to save memory, and specifying the compute type of <code>torch.float16</code> as that is the tensor type that this model was trained in. Using <code>torch.float32</code> (which is the default) will incur unnecessary conversion cost and slow down inference. <pre><code class="language-python">%pip install transformers~=4.36 torch~=2.1.1 bitsandbytes accelerate optimum from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig import torch quantization_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16 ) MODEL_NAME = "Upstage/SOLAR-10.7B-Instruct-v1.0" tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) model = AutoModelForCausalLM.from_pretrained( MODEL_NAME, device_map="auto", quantization_config=quantization_config, ) </code></pre> We define a generate function. <pre><code class="language-python">def generate(prompt: str, max_length: int = 2048): conversation = [{'role': 'user', 'content': prompt}] prompt = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True) inputs = tokenizer(prompt, return_tensors="pt").to(model.device) outputs = model.generate(**inputs, use_cache=True, max_length=max_length) output_text = tokenizer.decode(outputs[0]) return output_text </code></pre> And now we can test it out on a few prompts. Prompt: <code>Can you explain web3 to a 5 year old?</code> Response: Web3 is like a special world on the internet that lets people share and use things in a new and more open way. Imagine a playground where everyone can add their own swings, slides, and sandboxes without needing permission from a big boss. This makes it easier for people to work together and create new fun things online. (Inference time: <code>8s</code>) Prompt: <code>Can a firefighter transit into the finance industry?</code> Response: While it may not be a typical career transition, it is possible for a firefighter to transition into the finance industry. To make a successful move, the firefighter would need to acquire the necessary skills and qualifications relevant to the finance sector. This could involve pursuing additional education or training in finance, accounting, or business administration. Networking, gaining relevant work experience, and showcasing transferable skills from the firefighting profession can also help in making the transition.\n\nSome transferable skills from firefighting to finance could include strong problem-solving abilities, teamwork, decision-making under pressure, and the ability to manage risk. However, it's essential to note that the finance industry has its unique requirements, and the firefighter would need to demonstrate a strong commitment to learning and adapting to the new environment. (Inference time: <code>19s</code>) Seems to be pretty decent performance at a somewhat reasonable inference time out of the box. <a href="https://huggingface.co/docs/transformers/perf_infer_gpu_one">Speeding up GPU inference</a>. <h2 id="gemma-models"><a class="header" href="#gemma-models">Gemma Models</a></h2> References: <ul> <li><a href="https://huggingface.co/blog/gemma">Huggingface Reference</a></li> <li><a href="https://huggingface.co/blog/gemma-peft">Fine-tuning Gemma with PEFT</a></li> <li><a href="https://huggingface.co/google/gemma-7b">gemma-7b Model Card</a></li> </ul> Google released open source small models under the Gemma label. It comes in two sizes: 2B and 7B parameters, each with base (pretrained) and instruction-tuned versions. All have a context length of <code>8k tokens</code>. Gemma integration is supported with <code>transformers==4.38</code>. Note that user consent is needed to access Gemma models, so one needs to create a HuggingFace account and <a href="https://huggingface.co/docs/hub/security-tokens">create an access token</a>. <h3 id="sample-code"><a class="header" href="#sample-code">Sample Code</a></h3> <pre><code class="language-python"># Use model_name="google/gemma-7b-it" def load(model_name: str): tokenizer = AutoTokenizer.from_pretrained(model_name, token=TOKEN) model = AutoModelForCausalLM.from_pretrained( model_name, quantization_config=BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16, bnb_4bit_use_double_quant=True, # nested quantization bnb_4bit_quant_type="nf4", # normalized quantization ), token=TOKEN, ) return model, tokenizer </code></pre> <h3 id="inference-speed"><a class="header" href="#inference-speed">Inference Speed</a></h3> Tested out the inference speed of the model using the following prompt: <code>Can you recommend good alternative occupations for a occupation_name? Please be concise</code>. An example is as follows, for an occupation <code>counsellor</code>: <pre><code>Sure, here are some alternative occupations for a counsellor: * **Social worker:** Similar to counselling, social workers help individuals and families facing various challenges. * **Community organizer:** Organizes and leads community initiatives to address social issues. * **Teacher:** Can be a good fit for those who enjoy working with children and adolescents. * **Human resource manager:** HR managers handle employee relations and conflicts. * **Public relations specialist:** Helps organizations communicate effectively with the public. </code></pre> The inference speed is about <code>8-10 seconds</code> with 4-bit quantization. <h3 id="gemma-2b"><a class="header" href="#gemma-2b">Gemma-2b</a></h3> Testing out <code>google/gemma-2b-it</code> shows that it has poorer performance but faster inference speed at 2 billion parameters. Here is an example of the results for the same prompt. <pre><code>* **Life coach** * **Coach in education** * **Career coach** * **Executive coach** * **Leadership coach** * **Management consultant** * **Business development consultant** * **Financial advisor** * **Investment banker** </code></pre> The inference speed is as follows: <ul> <li>Single: <code>3s</code></li> <li>Batch of <code>4</code>: <code>14s</code></li> <li>Batch of <code>8</code>: <code>16s</code></li> <li>Batch of <code>16</code>: <code>20s</code></li> <li>Batch of <code>32</code>: OOM on T4 GPU</li> </ul> <h2 id="phi-2"><a class="header" href="#phi-2">Phi-2</a></h2> Mobius Labs fine-tuned the <code>phi-2</code> model from Microsoft which seems promising, and released it under <code>mobiuslabsgmbh/aanaphi2-v0.1</code>. The output seems better than <code>gemma-2b-it</code>. <pre><code>1. Social worker 2. Mental health therapist 3. School counselor 4. Employee assistance program (EAP) specialist 5. Rehabilitation counselor 6. Family therapist 7. Substance abuse counselor 8. Career counselor 9. Trauma-focused therapist 10. Child and adolescent therapist These occupations involve working with individuals, families, and communities to promote mental health and well-being, and may provide similar skills and experiences to those of a counsellor. </code></pre> <ul> <li><a href="https://colab.research.google.com/drive/1ge2F1QSK8Q7h0hn3YKuBCOAS0bK8E0wf#scrollTo=VPD7QS_DR-mw">bitsandbytes example notebook from HuggingFace</a></li> </ul> Flash attention, torch.compile, quantization. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="encoder-vs-decoder"><a class="header" href="#encoder-vs-decoder">Encoder vs Decoder</a></h1> There are broadly two categories of LLMs: Encoder-Decoder architecture (typified by BERT) and Decoder only architecture (typified by GPT-2 series). There are some innate differences between the two that affect the type of application each is well suited for. <h2 id="encoder-decoder-architecture"><a class="header" href="#encoder-decoder-architecture">Encoder-Decoder Architecture</a></h2> The Encoder-Decoder architecture was proposed in the original <a href="https://ar5iv.labs.arxiv.org/html/1706.03762">Attention is All You Need paper by Vaswani et al</a>. As the name suggests, there is a separate encoder and decoder module. The design choice of two distinct modules is easier to understand when we consider that the original paper was designed for the machine translation task, such that the inputs may be in English and the outputs may be in French. On the encoder-side, the encoder module sees the full input sequence and encodes it into an numeric encoding representation matrix of size <code>(seq_len, embed_dim)</code>. There is no input masking required because we assume we see the full English text before we begin translation. This encoding representation is passed to the decoder module as context. On the decoder-side, we cannot allow the decoder module to see the full output sequence during training, since that would cause data leakage and the model to collapse. Note that the decoder simultaneously (in one forward pass) predicts a token for each position in the output sequence during training. Hence to avoid data leakage, causal masking is applied such that for each position being predicted, the decoder module is only allowed to "see" tokens in previous positions. Note that this can be implemented simply by masking the attention matrix in each self-attention block in a triangular manner. The feedforward layers only map information from position <code>i</code> to position <code>i</code>, so there is no need to adjust that. <div class="table-wrapper"><table><thead><tr><th style="text-align: center"><img src="llm/../images/encoder_decoder_architecture.png" alt="BERT Encoder Decoder Architecture" /></th></tr></thead><tbody> <tr><td style="text-align: center">BERT Encoder Decoder Architecture (From Attention is All You Need)</td></tr> </tbody></table> </div> The decoder module itself creates a representation of the output sequence (with causal masking) of size <code>(seq_len, embed_dim)</code>. Now for the tricky part of merging the encoder and decoder representations. Naively, we can simply concatenate the two to form a representation of <code>(seq_len x 2, embed_dim)</code>, and add further self-attention blocks (with appropriate causal masking). However, this would increase the number of parameters of the model. The paper instead chose to take the <code>Q</code> matrix from the decoder and the <code>K, V</code> matrices from the encoder. Since the attention mechanism is softmax(Q⋅KT)V, this allows every token on the decoder representation to attend to every token on the encoder representation. The result is a weighted combination of the <code>V</code> vectors taken from the encoder representation which is eventually used to generate the tokens for each decoder position. Here we can already observe some inductive biases baked into the encoder-decoder architecture. These comments are inspired by / taken from <a href="https://www.youtube.com/watch?v=orDKvo8h71o">Hyung Won Chung's lecture</a>: <ol> <li>There is a clear distinction between the tokens on the encoder-side and the decoder-side, such that tokens on either side cannot directly attend to each other until the final block when the mixing happens. This may be suitable for the translation task but not for generic chat completion, where it is quite arbitary which point in the conversation we determine as the encoding and which point we determine as the decoding.</li> <li>The <code>V</code> matrix is taken entirely from the encoder-side, which may again limit the expressiveness of the model (what if there are important tokens on the decoder-side that would be useful to taken the <code>V</code> values from?). One may argue that in translation, the tokens on the encoder-side capture the full meaning of the sentence to be translated, so it is comprehensive. But the same is not true of generic chat completions.</li> <li>The separation of an encoder and decoder module suggests that there is significant difference between the tokens on the encoder-side and decoder-side. Again, this inductive bias</li> </ol> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="contextualized-recommendations"><a class="header" href="#contextualized-recommendations">Contextualized Recommendations</a></h1> Contextualized recommendations is an emerging use case from LLMs. The idea is that we use a traditional search and recommender system to generate recommendations, and use an LLM to craft an explanation for why each recommendation is relevant to the user in a very personalized way. Multiple companies have reported that this has driven engagement and clicks up significantly. <h2 id="spotify"><a class="header" href="#spotify">Spotify</a></h2> <a href="https://research.atspotify.com/2024/12/contextualized-recommendations-through-personalized-narratives-using-llms/">Contextualized recommendations through personalized narratives using LLMs</a> Traditionally, spotify users use just the cover art to decide whether to engage with a new music recommendation. Spotify wants to include a short one-liner to explain why a particular item might resonate with users. For example, <code>Dead Rabbitts latest single is a metalcore adrenaline rush!</code> or <code>Relive U2's iconic 1993 Dublin concert with ZOO TV Live EP</code>. Spotify highlights some challenges they faced: <ul> <li>Ensuring a consistent generation style and tone</li> <li>Avoiding harmful or inappropriate outputs</li> <li>Mitigating hallucinations and inaccuracies</li> <li>Understanding user preferences to deliver tailored meaningful explanations</li> </ul> Initial tests with zero-shot / few-shot Llama did not work too well. They adopted a human-in-the-loop approach: <ul> <li>Expert editors provide "golden examples" for instruction fine-tuning</li> <li>Provide ongoing feedback to address errors in LLM output <ul> <li>Artist attribution errors</li> <li>Tone inconsistencies</li> <li>Factual inaccuracies</li> </ul> </li> </ul> The AB tests showed that explanations containing meaningful details about the artist or music led to significantly higher user engagement. For LLM fine-tuning, they found that Llama 3.1 8B worked well and could be trained with multiple adapters for 10 different tasks. Throughout the training process, they used MMLU benchmark as a guardrail to ensure that the model's overall ability remained intact. Spotify uses vLLM for inference. <h2 id="linkedin"><a class="header" href="#linkedin">LinkedIn</a></h2> <a href="https://www.linkedin.com/pulse/celebrating-1-billion-members-our-new-ai-powered-linkedin-tomer-cohen-26vre">Our new AI powered LinkedIn</a> LinkedIn provides AI features for premium users. When users click on a job, they can ask questions like "Am I a good fit for the job?". The LLM will respond with a short bullet-pointed explanation on: <ul> <li>Whether the user is a good fit</li> <li>Details from the user's profile that make them a good fit</li> <li>Areas that the users are missing</li> </ul> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="mingpt"><a class="header" href="#mingpt">MinGPT</a></h1> This is a walkthrough of KArpathy's MinGPT implementation <h2 id="bpepy"><a class="header" href="#bpepy"><code>bpe.py</code></a></h2> This file contains code to implement a Byte Pair Encoding encoder. It does not contain code for training, just loads openai's GPT-2 bpe encoding for inference. Most text is represented in UTF-8 encoding, which is just a sequence of bytes (values <code>0 to 255</code>). For example, <code>0xf0</code> corresponds to <code>33</code> in decimal which corresponds to the character <code>!</code>. This means that all text can be treated as a sequence of byte values. As a fallback, we first need token representations for individual byte values (in case we encounter out of vocab tokens with unknown byte sequences). First, two files are downloaded in <code>get_encoder</code>, from <code>https://openaipublic.blob.core.windows.net/gpt-2/models/124M/</code>: <ul> <li><code>encoder.json</code> is a dict of len <code>50257</code> mapping from a token to its index. This represents the entirety of the vocabulary <ul> <li>The first <code>256</code> tokens represents the <code>256</code> byte values. Each token is some arbitrarily chosen character (just need to make sure it is printable)</li> <li>These tokens are a fallback to ensure that we can encode any text sequence. For example, if we encounter a new emoji with unknown byte sequence, at the very least we can encode each byte separately.</li> <li>The next <code>50k</code> tokens map from a byte sequence of length <code>2</code> and above to an index, these are the BPE mined sequences of merged bytes</li> <li>The last token is <code><|endoftext|></code> which is a special token</li> </ul> </li> <li><code>vocab.bpe</code> is a <code>\n</code> separated list of byte sequences that should be merged (<code>50k</code> of them) <ul> <li>In contrast to the above, these sequences are not merged yet (e.g. a line is <code>R ocket</code>)</li> <li>We store these as a <code>list[tuple]</code> in <code>bpe_merges</code></li> </ul> </li> </ul> These two data are passed into the <code>Encoder</code> main class. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="statistics"><a class="header" href="#statistics">Statistics</a></h1> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="conformal-predictions"><a class="header" href="#conformal-predictions">Conformal Predictions</a></h1> <h2 id="exchangeability"><a class="header" href="#exchangeability">Exchangeability</a></h2> In this section we discuss the difference between independence and exchangeability. <blockquote> Definition. Independence. Two events A and B from a sample space S are said to be independent if P(A and B)=P(A)P(B). </blockquote> If P(B)>0, it is equivalent to say P(A∣B)=P(A), since: P(A∣B)=P(B)P(A,B)=P(B)P(A)P(B)=P(A) <blockquote> Definition. Exchangeability Two events A and B are said to be exchangeable if P(A1=a,A2=b)=P(A1=b,A2=a), which means there is indifference with respect to the order of events. More generally, exchangeability of a sequence of events A1,...,An means that the joint distribution is unchanged when we permute the order of events: P(A1,...,An)=P(Aσ(1),...,Aσ(n)), for all permutations σ </blockquote> The simplest way to understand is to use the example of drawing balls from an urn without replacement (example taken from Cordani 2006). Suppose we have an urn with <code>10 red balls</code> and <code>5 white balls</code>. Then the following tree shows the draw probabilities at each step: <pre class="mermaid">graph TD; A((Start)) --> R1("R1 (10/15)"); A --> W1("W1 (5/15)"); R1 --> R2("R2 (9/14)"); R1 --> W2("W2 (5/14)"); W1 --> R2_2("R2 (10/14)"); W1 --> W2_2("W2 (4/14)"); R2 --> R3("R3 (8/13)"); R2 --> W3("W3 (5/13)"); W2 --> R3_2("R3 (9/13)"); W2 --> W3_2("W3 (4/13)"); R2_2 --> R3_3("R3 (9/13)"); R2_2 --> W3_3("W3 (4/13)"); W2_2 --> R3_4("R3 (10/13)"); W2_2 --> W3_4("W3 (3/13)"); </pre> Suppose we <h2 id="references-3"><a class="header" href="#references-3">References</a></h2> <ul> <li><a href="https://iase-web.org/documents/papers/icots7/3I1_CORD.pdf">Cordani 2006 - TEACHING INDEPENDENCE AND EXCHANGEABILITY</a></li> <li><a href="https://www.stat.berkeley.edu/~ryantibs/statlearn-s23/lectures/conformal.pdf">Tibshirani 2023 - Conformal Predictions</a></li> </ul> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="miscellaneous-notes"><a class="header" href="#miscellaneous-notes">Miscellaneous Notes</a></h1> A collection of miscellaneous, useful notes. <h2 id="f-strings"><a class="header" href="#f-strings">f-strings</a></h2> To surround a <code>text</code> with a symbol (say <code>=</code>) to a fixed length: <pre><code class="language-python">>>> text = "Title" >>> print(f"{text:=^20}") =======Title======== </code></pre> <h2 id="vim"><a class="header" href="#vim">Vim</a></h2> Command to interactively change each 'foo' to 'bar'. <code>:%s</code> triggers the substitute global command, followed by the search and replace phrases respectively. Finally <code>g</code> means replace all occurrences and <code>c</code> means with confirmation. Note that <code>:s</code> will only do the same for one line. <pre><code class="language-bash">:%s/foo/bar/gc </code></pre> <h2 id="find-files"><a class="header" href="#find-files">Find Files</a></h2> To find files anywhere on the system with the filename <code>python</code> using bash, use: <pre><code class="language-bash">find . -name python </code></pre> We can add <code>*</code> before and/or after the filename to allow other characters before or after our keyword: <pre><code class="language-bash">find . -name *python* </code></pre> To search not just in the filename but also in the full path (e.g. we only want to search in <code>Desktop</code>), we can do: <pre><code class="language-bash">find . -wholename "*Desktop*python*" </code></pre> Note that if we want to locate executable binaries, another useful command is <code>whereis</code>: <pre><code class="language-bash">whereis cat --- cat: /usr/bin/cat /usr/share/man/man1/cat.1.gz </code></pre> <h2 id="numpy-indexing"><a class="header" href="#numpy-indexing">Numpy Indexing</a></h2> Suppose we have a 2D array <code>X</code> and would like to take a slice of certain rows and columns. We might try, for e.g., to take the first two rows of X and the 3rd/4th column of X, i.e. we expect to get a 2 by 2 matrix. <pre><code class="language-python">import numpy as np X = np.random.randn(5, 5) X[[0, 1], [2, 3]] </code></pre> Unfortunately, this will return an array of items <code>array(X[0, 2], X[1, 3])</code>, which is not what we want. Instead, a slightly inefficient but clear way is to slice each axis separately. It is not optimal because the first slice creates a temporary array view before the second slice is applied. <pre><code class="language-python">X[[0, 1], :][:, [2, 3]] </code></pre> Finally, the recommended way seems to be to use <code>np.ix_</code>. Doing <code>np.ix_([0, 1], [2, 3])</code> creates a tuple of two elements. <pre><code class="language-python">idx = np.ix_([0, 1], [2, 3]) idx >> (array([[0], [1]]), >> array([[2, 3]])) </code></pre> Indexing the original array with this output <code>X[idx]</code> will then give us what we want. <h2 id="asyncio"><a class="header" href="#asyncio">asyncio</a></h2> <a href="https://realpython.com/async-io-python/">reference</a> asyncio is a single-threaded framework that does not use multi-threading or multi-processing to speed up tasks. Instead, a coordinator (or event loop) passes control from a time-consuming blocking function (e.g. time.sleep or an I/O operation) to other functions to run. This passing of control occurs with the <code>await</code> keyword. When the blocking function is completed, it notifies the coordinator and control returns to where the blocking function left off. <code>asyncio</code> does not speed up CPU-bound tasks, due to its single-threaded design. It only works when the function being awaited is an I/O operation that is supported by asyncio. This includes <a href="https://realpython.com/async-io-python/#libraries-that-work-with-asyncawait">stuff</a> like: <ul> <li>HTTP (supported through <code>aiohttp</code>)</li> <li>DB calls (e.g. <code>aioredis</code>)</li> </ul> <code>asyncio</code> is preferred for such tasks (e.g. making multiple I/O calls and doing something when they all return) over multi-processing because of its single-thread design, making it easier to debug. <h2 id="relatedness"><a class="header" href="#relatedness">relatedness</a></h2> In the context of information retrieval, Trey Grainger in <a href="https://www.manning.com/books/ai-powered-search">AI-Powered Search</a> suggests a relatedness measure to connect arbitrary entities together. Suppose we have a collection of jobs and each job is tagged with a set of skills. Suppose we wish to retrieve relevant skills to an arbitrary free text query q. The relatedness idea is to define a foreground of documents, e.g. based on a retrieval of documents using query q which are related to the query, and to compare the attributes of the foreground against the background, i.e. all documents. Mathematically, we can think of the foreground documents as a sample, and the background documents as the population. The strength of the relationship between each skill t to the query q may then be defined as the z-statistic of the one-sample z-test of proportions of the occurrence of skill t in the foreground sample compared against the background population. A significantly greater occurrence in the sample compared to the population suggests a strong relationship between t and q, and vice versa. Specifically: z=np(1−p)<svg xmlns="http://www.w3.org/2000/svg" width='400em' height='1.88em' viewBox='0 0 400000 1944' preserveAspectRatio='xMinYMin slice'><path d='M983 90 l0 -0 c4,-6.7,10,-10,18,-10 H400000v40 H1013.1s-83.4,268,-264.1,840c-180.7,572,-277,876.3,-289,913c-4.7,4.7,-12.7,7,-24,7 s-12,0,-12,0c-1.3,-3.3,-3.7,-11.7,-7,-25c-35.3,-125.3,-106.7,-373.3,-214,-744 c-10,12,-21,25,-33,39s-32,39,-32,39c-6,-5.3,-15,-14,-27,-26s25,-30,25,-30 c26.7,-32.7,52,-63,76,-91s52,-60,52,-60s208,722,208,722 c56,-175.3,126.3,-397.3,211,-666c84.7,-268.7,153.8,-488.2,207.5,-658.5 c53.7,-170.3,84.5,-266.8,92.5,-289.5z M1001 80h400000v40h-400000z'/></svg>p^−p Where: <ul> <li>p^=df(q)df(q,t) is the sample proportion.</li> <li>df(q,t) is the number of documents in the foreground corresponding to query q and contains skill t.</li> <li>df(q) is the total number of documents in the foreground corresponding to query q. It is also the number of samples n.</li> <li>p=P(t) is the probability of skill t appearing across all documents.</li> </ul> By performing a retrieval and ranking skills based on the z-statistic, we can compute a relationship between any arbitrary query and attribute of the documents (on the fly, if necessary). This functionality is implemented in <a href="https://solr.apache.org/guide/8_5/json-facet-api.html#relatedness-and-semantic-knowledge-graphs">solr</a>. <h2 id="package-versioning"><a class="header" href="#package-versioning">Package Versioning</a></h2> <a href="https://python-poetry.org/docs/dependency-specification/">link</a> <a href="https://semver.org/">SemVer</a> is the versioning standard for Python packages. For a package version of e.g. <code>1.2.3</code>: <ul> <li>The first number <code>1</code> is the major version. We update it when we make major API changes. A major of <code>0</code> indicates that the package is under initial development and anything may change.</li> <li>The second number <code>2</code> is the minor version. We update it when we add functionality in a backward compatible manner.</li> <li>The third number <code>3</code> is the patch version. We update it when we patch bugs in a backward compatible manner.</li> </ul> Poetry is a useful tool to manage package dependencies in a Python library. In the <code>pyproject.toml</code> file, we specify package dependencies in the following manner: <pre><code>[tool.poetry.dependencies] python = ">=3.8,<3.11" requests = "^2.26.0" pytz = "~2022.1" </code></pre> The <code>caret</code> requirement (e.g. <code>requests = "^2.26.0"</code>) means that SemVer compatible changes are allowed, i.e. an update is allowed if the version number does not modify the left-most non-zero digit. E.g.: <ul> <li><code>^1.2.3</code> means that <code>1.3.0</code> is allowed but not <code>2.0.0</code></li> <li><code>^0.2.4</code> means that <code>0.2.5</code> is allowed but not <code>0.3.0</code></li> </ul> The <code>tilde</code> requirement is somewhat stricter. It specifies a minimal version but allows room to upgrade depending on how many digits are supplied. If <code>major.minor.patch</code> or <code>major.minor</code> is supplied, only patch-level changes are allowed. If <code>major</code> is supplied, minor and patch level changes are allowed. <ul> <li><code>~1.2.3</code> means that <code>1.2.7</code> is allowed but not <code>1.3.0</code></li> <li><code>~1.2</code> means that <code>1.2.7</code> is allowed but not <code>1.3.0</code></li> <li><code>~1</code> means that <code>1.3.0</code> is allowed but not <code>2.0.0</code></li> </ul> The poetry versioning specification depends on developer updating the library versions in a disciplined way. It allows us to provide some flexibility in package versioning while avoiding updating a dependent package to a version that breaks our code. Having more flexibility in specifying package dependencies reduces the risk of dependencies version conflicting. <h2 id="parallelism-in-python"><a class="header" href="#parallelism-in-python">Parallelism in Python</a></h2> References: <a href="https://realpython.com/python-gil/">realpython</a>, <a href="https://superfastpython.com/numpy-vs-gil/">superfastpython</a> The Python Global Interpreter Lock (GIL) is a lock that allows only one thread to run at a time. The reason for this design feature in Python is because the Python interpreter is not thread-safe, meaning that if the GIL did not exist, the Python interpreter could introduce race conditions and other fatal errors due to multiple threads modifying memory at the same time. The GIL also allowed other non-thread-safe C programs to be easily integrated into the Python ecosystem, which contributed to Python's success as a language. The GIL also improves performance for single-threaded programs, as it only requires a single lock to be managed (not sure how this works). There have been attempts to remove the GIL from Python but none have been successful because they would degrade single-threaded performance. For many C extensions (e.g. numpy), multi-threading is still possible because these extensions are allowed to manually release the GIL (as long as they restore things back to normal when the functions return). This allows us to still use multi-threading for CPU-intensive functions with the GIL around. Similarly for Rust, we can release the GIL to achieve <a href="https://pyo3.rs/v0.9.2/parallelism">parallelism</a>. Alternatively, we can use <code>multiprocessing</code> to create multiple processes (instead of threads). Each process contains its own Python interpreter (and GIL) and hence can run truly in parallel. The downside is that the overhead of creating and managing processes is much more than that for threads, meaning that the benefits of multiprocessing are much dampened compared to multi-threading. <h2 id="memory-profiling"><a class="header" href="#memory-profiling">Memory Profiling</a></h2> It is often useful to profile the memory usage of our script. In python, we can use <code>memory_profiler</code> to check the memory usage of our program line by line. <pre><code class="language-python">from memory_profiler import profile import sys @profile(stream=sys.stdout) def f(): a = [1] * (10**6) b = [2] * (2 * 10**7) del b if __name__ == "__main__": f() </code></pre> This will print the following useful information to stdout. Note that even before we did anything, there is background memory usage of 17MiB. <pre><code>Filename: memory_profiling.py Line # Mem usage Increment Occurrences Line Contents ============================================================= 5 17.1 MiB 17.1 MiB 1 @profile(stream=sys.stdout) 6 def f(): 7 24.5 MiB 7.5 MiB 1 a = [1] * (10**6) 8 177.2 MiB 152.6 MiB 1 b = [2] * (2 * 10**7) 9 24.8 MiB -152.4 MiB 1 del b </code></pre> We might also want to track memory usage of a function over time. We can use <code>memory_usage</code> instead for that. <pre><code class="language-python">import time from memory_profiler import memory_usage def g(a, n: int = 100): time.sleep(1) b = [a] * n time.sleep(1) del b time.sleep(1) if __name__ == "__main__": usage = memory_usage((g, (1,), {"n": int(1e7)}), interval=0.5) print(usage) </code></pre> This will give us an array like so, showing the spike in memory in the middle of g. <pre><code>[17.375, 17.5234375, 17.5234375, 19.34765625, 93.59765625, 93.59765625, 17.53125, 17.53125, 17.53125] </code></pre> <h2 id="cuda"><a class="header" href="#cuda">CUDA</a></h2> PyTorch may sometimes throw errors if the installed torch version does not match the installed CUDA version. To address this, we need to first check the CUDA version using the <code>nvcc</code> command: <pre><code class="language-bash">/usr/local/cuda/bin/nvcc --version --- nvcc: NVIDIA (R) Cuda compiler driver Copyright (c) 2005-2022 NVIDIA Corporation Built on Wed_Jun__8_16:49:14_PDT_2022 Cuda compilation tools, release 11.7, V11.7.99 Build cuda_11.7.r11.7/compiler.31442593_0 </code></pre> Then install the correct version of torch: <pre><code class="language-bash">pip install -U torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 </code></pre> <h2 id="list-comprehension"><a class="header" href="#list-comprehension">List Comprehension</a></h2> When doing nested list comprehensions in python, the later loop becomes nested as the inner loop. So this: <pre><code class="language-python">l = [i * j for i in range(4) for j in range(5)] </code></pre> Is equivalent to: <pre><code class="language-python">l = [] for i in range(4): for j in range(5): l.append(i*j) </code></pre> This also explains why flattening a list of lists is of the following form: <pre><code class="language-python">list_of_lists = [[0, 1, 2], [3, 4, 5], [6, 7, 8]] flattened = [item for sublist in list_of_lists for item in sublist] </code></pre> When we translate it into a nested for loop, we can see that the inner loop should indeed be placed at the end: <pre><code class="language-python">flattened = [] for sublist in list_of_lists: for item in sublist: flattened.append(item) </code></pre> <h2 id="configure-vscode-snippets"><a class="header" href="#configure-vscode-snippets">Configure VsCode Snippets</a></h2> When typing stuff in VsCode, we sometimes write the same boilerplate and want to have a shortcut to type it for us. For example, when writing latex, we often have to create an <code>align</code> environment like so: <pre><code class="language-bash">$$ \begin{align*} Equations go here.. \end{align*} $$ </code></pre> VsCode has a snippets function to help us with that. <ol> <li>Press <code>Ctrl + Shift + P</code> and type <code>Configure User Snippets</code></li> <li>Create a local or global snippets file and name it (e.g. <code>charles.code-snippets</code>)</li> <li>Add the snippet definition like the below</li> </ol> <pre><code class="language-bash">{ "LaTeX Align Environment": { "prefix": "align", "body": [ "$$", "\\begin{align*}", "\t$1", "\\end{align*}", "$$$0" ], "description": "Inserts a LaTeX align* environment." } } </code></pre> Note that the prefix is the keyword to trigger the template. To use the template, we simply enter <code>insert</code> mode in vim as per normal (e.g. using <code>a</code> or <code>i</code> or <code>o</code>), then type <code>align</code>. At this point, we would see the text as per normal. We can then press <code>ctrl+space</code> to show the snippets auto-completion, if there are any. Finally, we just press <code>enter</code> to get the template generated. The cursor will now come to rest at position <code>$1</code> as indicated in the snippet. <h2 id="github-in-colab"><a class="header" href="#github-in-colab">Github in Colab</a></h2> We may want to run <code>.ipynb</code> notebooks within a github project in colab. There are two steps to this. First, to get the <code>.ipynb</code> itself into colab, we can either use the UI to manually upload the notebook, or if we are loading from Github, we can simply start with the url to the notebook <code>https://github.com/charleslow/REPOSITORY/blob/master/NOTEBOOK.ipynb</code> and change <code>github</code> to <code>githubtocolab</code>. This will automatically open up the notebook in colab. However, this only gives us the standalone notebook, and we will not have access to the project directory and packages. Hence, in the notebook, we need to: <ul> <li><code>!git clone <Insert REPOSITORY name>.git</code></li> <li><code>%cd folder name</code></li> </ul> This will allow us to import packages and run the notebook as per normal. <h2 id="quartz-cron"><a class="header" href="#quartz-cron">Quartz Cron</a></h2> We often use quartz cron expressions to schedule jobs. A cron expression is simply 6 or 7 values in a space-separated string, such as: <pre><code>"0 0 0 1 * ?" </code></pre> Which means run on the: <ul> <li><code>0th second</code> of the</li> <li><code>0th minute</code> of the</li> <li><code>0th hour</code> of the</li> <li><code>1st day the month</code> of</li> <li><code>every month</code></li> <li>and doesn't matter which <code>day of the week</code></li> </ul> The specifics of each field are in the table below, from <a href="https://www.quartz-scheduler.org/documentation/quartz-2.3.0/tutorials/crontrigger.html">the cron trigger tutorial</a>: <div class="table-wrapper"><table><thead><tr><th style="text-align: left">Field Name</th><th style="text-align: left">Mandatory</th><th style="text-align: left">Allowed Values</th><th style="text-align: left">Allowed Special Characters</th></tr></thead><tbody> <tr><td style="text-align: left">Seconds</td><td style="text-align: left">YES</td><td style="text-align: left">0-59</td><td style="text-align: left"><code>, - * /</code></td></tr> <tr><td style="text-align: left">Minutes</td><td style="text-align: left">YES</td><td style="text-align: left">0-59</td><td style="text-align: left"><code>, - * /</code></td></tr> <tr><td style="text-align: left">Hours</td><td style="text-align: left">YES</td><td style="text-align: left">0-23</td><td style="text-align: left"><code>, - * /</code></td></tr> <tr><td style="text-align: left">Day of month</td><td style="text-align: left">YES</td><td style="text-align: left">1-31</td><td style="text-align: left"><code>, - * ? / L W</code></td></tr> <tr><td style="text-align: left">Month</td><td style="text-align: left">YES</td><td style="text-align: left">1-12 or JAN-DEC</td><td style="text-align: left"><code>, - * /</code></td></tr> <tr><td style="text-align: left">Day of week</td><td style="text-align: left">YES</td><td style="text-align: left">1-7 or SUN-SAT</td><td style="text-align: left"><code>, - * ? / L #</code></td></tr> <tr><td style="text-align: left">Year</td><td style="text-align: left">NO</td><td style="text-align: left">empty, 1970-2099</td><td style="text-align: left"><code>, - * /</code></td></tr> </tbody></table> </div> On the special characters: <ul> <li><code>*</code> is commonly used, meaning "select all the values"</li> <li><code>?</code> is only used for Day of month or Day of week. Basically if we set one, we must set the other to <code>?</code> to avoid a conflict</li> <li><code>-</code> is used to specify a range. E.g. <code>10-12</code> in the hour field means run on the <code>10th</code>, <code>11th</code> and <code>12th</code> hour</li> <li><code>,</code> is used to specify a few values, e.g. <code>10,12</code> n the hour field means run on the <code>10th</code> and <code>12th</code> hour</li> <li><code>/</code> is used to specify increments. e.g. <code>0/15</code> in the minutes field means run on minutes <code>0, 15, 30, 45</code></li> <li><code>#</code> is a special case for Day of the Week. Basically, <code>MON#1</code> means the 1st occurrence of a Monday of the given month</li> </ul> That's about it! Or we can just use an LLM to generate the expression we need. <h2 id="wayland-in-devcontainer"><a class="header" href="#wayland-in-devcontainer">Wayland in Devcontainer</a></h2> Sometimes we may encounter this error when trying to use vscode to setup a devcontainer: <pre><code>Error response from daemon: invalid mount config for type "bind": bind source path does not exist: /run/user/1000/wayland-0 </code></pre> This happens because VS Code is trying to share the graphics connection with the container, which is not needed in most cases. Here is how to disable it: <ul> <li>Open VS Code Settings <code>(Ctrl + ,)</code></li> <li>Search for "Mount Wayland Socket"</li> <li>Uncheck the box for Dev > Containers: Mount Wayland Socket</li> <li>Rebuild your container (Ctrl + Shift + P -> "Rebuild Container").</li> </ul> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="bradley-terry-model"><a class="header" href="#bradley-terry-model">Bradley-Terry Model</a></h1> Based on <a href="https://en.wikipedia.org/wiki/Bradley%E2%80%93Terry_model">wikipedia</a>. The Bradley-Terry model (1952) is a probability model that allows us to infer scores for individual objects based on a dataset of pairwise comparisons between them. Specifically, it estimates the probability that i▹j (i.e. i is preferred to j) as: P(i▹j)=pi+pjpi Where pi is a positive real-valued score assigned to object i (not necessarily a probability). Typically, pi is parametrized as an exponential score pi=eβi, and the goal is to learn the parameters βi from pairwise comparisons. This results in: P(i▹j)=eβi+eβjeβi <h2 id="parameter-estimation"><a class="header" href="#parameter-estimation">Parameter Estimation</a></h2> Parameter estimation is typically done using maximum likelihood. Starting with a set of pairwise comparisons between individual objects, let wij be the number of times object i beats object j. Then the likelihood of a given set of parameters p:=[p1,...,pn] (n denotes number of objects) is as follows: L(p)=lnij∏P(i▹j)wij=i=1∑nj=1∑nln(pi+pjpi)wij=i=1∑nj=1∑nwij[lnpi−ln(pi+pj)] This likelihood function can then be minimized by differentiating wrt pi and solved by setting to zero. <h2 id="bt-model-as-a-sigmoid-function"><a class="header" href="#bt-model-as-a-sigmoid-function">BT Model as a Sigmoid Function</a></h2> We can also express the likelihood as a function of the difference in scores βi−βj. Recall that the sigmoid function is σ(x)=1/(1+e−x). Then: L(p)=i=1∑nj=1∑nwij⋅ln[eβi+eβjeβi]=i=1∑nj=1∑nwij⋅ln σ(βi−βj) The derivation for the second line above is found at <a href="misc/../identities/sigmoid.html">/Identities/sigmoid</a>. This re-parametrization shows that the BT-model is really modelling the preference as a difference in scores and then running that through the sigmoid function to convert it into a probability. This means that we are basically running a pairwise logistic regression. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="setting-up-wsl"><a class="header" href="#setting-up-wsl">Setting up WSL</a></h1> We may need to uninstall and re-install WSL (Windows Subsystem for Linux) from time to time. Here is the step-by-step. <ol> <li>Tear down and delete all files. <code>wsl --unregister Ubuntu-22.04</code></li> <li>Re-install wsl. <code>wsl --install Ubuntu-22.04</code></li> <li>Set up keychain to auto-find ssh-agent (from Windows) and add keys <ul> <li>Copy <code>.ssh</code> folder from Windows to wsl <code>~/.ssh</code></li> <li>Add to <code>~/.bashrc</code> the following: <code>eval $(keychain --eval --agents ssh id_rsa)</code></li> </ul> </li> <li>Install Docker <ul> <li>Add users to docker to allow vscode access</li> <li>https://docs.docker.com/engine/install/linux-postinstall/</li> </ul> </li> <li>Essentials (get C linker) <ul> <li><code>sudo apt install build-essential</code></li> </ul> </li> <li>Install rust <ul> <li>https://www.rust-lang.org/tools/install</li> </ul> </li> <li>Install mdbook <ul> <li><code>cargo install mdbook mdbook-katex</code></li> </ul> </li> </ol> <h1 id="ssh-agent-forwarding"><a class="header" href="#ssh-agent-forwarding">ssh-agent forwarding</a></h1> One challenge in WSL is to forward the ssh-agent from WSL into the devcontainer. Make sure to include the following in <code>devcontainer.json</code> for this to happen. <pre><code class="language-json">{ "customizations": { "vscode": { "settings": { "remote.containers.sshForwardAgent": true } } } } </code></pre> <h2 id="configure-blackout"><a class="header" href="#configure-blackout">Configure Blackout</a></h2> It's useful to have a screen blackout command to keep jobs running but not show our screen when leaving the desk. Click <code>New</code> and create a shortcut, paste the following command: <pre><code class="language-powershell">powershell.exe -Command "(Add-Type '[DllImport(\"user32.dll\")]public static extern int SendMessage(int hWnd, int hMsg, int wParam, int lParam);' -Name a -Pas)::SendMessage(-1,0x0112,0xF170,2)" </code></pre> Set the shortcut name (e.g. <code>blackout</code>). Finally, add a keyboard shortcut to it (e.g. <code>Ctrl + Alt + B</code>). <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="to-read"><a class="header" href="#to-read">To Read</a></h1> <ul> <li>https://netflixtechblog.com/improve-your-next-experiment-by-learning-better-proxy-metrics-from-past-experiments-64c786c2a3ac</li> <li>https://netflixtechblog.com/sequential-a-b-testing-keeps-the-world-streaming-netflix-part-1-continuous-data-cba6c7ed49df</li> <li>https://arxiv.org/abs/2407.02464</li> <li>https://arxiv.org/abs/2305.12102</li> <li>https://arxiv.org/abs/1709.03933</li> <li>https://arxiv.org/abs/2101.08769</li> <li>https://arxiv.org/abs/1205.2618 # BPR</li> <li>https://arxiv.org/abs/2312.08520 # InfoNCE+</li> <li>https://arxiv.org/abs/2308.06091 # MAWU</li> <li>https://arxiv.org/abs/2109.12613 # CCL</li> <li>https://arxiv.org/abs/2201.02327 # SSM</li> <li>https://arxiv.org/abs/2206.12811 # DirectAU</li> <li><a href="https://arxiv.org/abs/1909.03193">KG-BERT</a></li> <li><a href="https://arxiv.org/abs/2110.00685">XR-Transformer</a></li> <li><a href="https://arxiv.org/abs/2209.11055">SetFit</a></li> <li><a href="https://arxiv.org/abs/2410.21228">Lora vs Fine tuning</a></li> <li><a href="https://arxiv.org/abs/2306.08121">Better Generalization with semantic IDs</a></li> <li><a href="https://arxiv.org/abs/1801.04062">Mutual Information Neural Estimation</a></li> <li><a href="https://arxiv.org/abs/2010.04592">Contrastive Learning with Hard Negatives</a></li> <li><a href="https://arxiv.org/abs/2110.06848">Decoupled Contrastive Learning</a></li> <li><a href="https://arxiv.org/abs/2202.06200">Improving Graph Collaborative Filtering with Neighborhood-enriched Contrastive Learning</a></li> <li><a href="https://arxiv.org/abs/1609.02907">Semi-Supervised Classification with Graph Convolutional Networks</a></li> <li><a href="https://arxiv.org/abs/1809.07428">Ranking Distillation</a></li> <li><a href="https://arxiv.org/abs/2205.04507">PinnerFormer</a></li> <li><a href="https://arxiv.org/abs/1904.06690">BERT4Rec</a></li> <li><a href="https://arxiv.org/abs/2212.02027">Retrieval as Attention</a></li> <li><a href="https://arxiv.org/abs/2209.10063">Generate Rather than Retrieve</a></li> <li><a href="https://arxiv.org/abs/2105.11108">PLM in Baidu Search</a></li> <li><a href="https://arxiv.org/abs/2006.11632">Emebdding Based Retrieval in Facebook Search</a></li> <li><a href="https://arxiv.org/abs/1808.09781">Self Attentive Sequential Recommendation</a></li> <li><a href="https://arxiv.org/abs/2309.07602">Is BERT4Rec really better than SASRec?</a></li> <li><a href="https://scontent.fsin4-1.fna.fbcdn.net/v/t39.8562-6/246721374_422204999475172_9039387325224382577_n.pdf?_nc_cat=104&ccb=1-7&_nc_sid=e280be&_nc_ohc=adv7Euc5H4IQ7kNvgEHyaVc&_nc_oc=Adh-DQ1sU1WSDRjqnODhkqDrDss2u4MkbNEHo74Sm7EaF2U0pPfNK_pqJf65fg1KKU0&_nc_zt=14&_nc_ht=scontent.fsin4-1.fna&_nc_gid=An-5fgAThp16mU1IsiAhSVP&oh=00_AYDbJULpLpX5Ekq9BeE8d2Ofno_1Boxkw3y_JnEdtA0RWA&oe=67B90BFA">Transformers4Rec</a></li> <li><a href="https://arxiv.org/abs/2010.06467">Pretrained Transformers for Text Ranking: BERT and Beyond</a></li> <li><a href="https://arxiv.org/abs/2309.07602">Turning Dross Into Gold Loss: is BERT4Rec really better than SASRec?</a></li> <li><a href="https://arxiv.org/abs/2004.13313">Modularized Transfomer-based Ranking Framework</a></li> <li><a href="https://arxiv.org/abs/2106.03373">Pre-trained Language Model for Web-scale Retrieval in Baidu Search</a></li> <li><a href="https://arxiv.org/abs/1702.07969">Related Pins at Pinterest: The Evolution of a Real-World Recommender System</a></li> <li><a href="https://arxiv.org/abs/2410.17152">Improving Pinterest Search Relevance Using Large Language Models</a></li> <li><a href="https://arxiv.org/abs/1711.07601">Pixie: A System for Recommending 3+ Billion Items to 200+ Million Users in Real-Time</a></li> <li><a href="https://arxiv.org/abs/2006.02282">Towards Personalized and Semantic Retrieval: An End-to-End Solution for E-commerce Search via Embedding Learning</a></li> <li><a href="https://www.kdd.org/kdd2016/papers/files/adf0361-yinA.pdf">Ranking Relevance in Yahoo Search</a></li> <li><a href="https://arxiv.org/abs/2310.08319">Fine-Tuning LLaMA for Multi-Stage Text Retrieval</a></li> <li><a href="https://arxiv.org/abs/2106.11959">Revisiting Deep Learning Models for Tabular Data</a></li> <li><a href="https://arxiv.org/abs/2404.16260">OmniSearchSage: Multi-Task Multi-Entity Embeddings for Pinterest Search</a></li> <li><a href="https://arxiv.org/abs/2006.05639">Search-based User Interest Modeling with Lifelong Sequential Behavior Data for Click-Through Rate Prediction</a></li> <li><a href="https://arxiv.org/abs/1905.06874">Behavior Sequence Transformer for E-commerce Recommendation in Alibaba</a></li> <li><a href="https://arxiv.org/abs/1806.01973">Graph Convolutional Neural Networks for Web-Scale Recommender Systems</a></li> <li><a href="https://arxiv.org/abs/2007.03634">PinnerSage: Multi-Modal User Embedding Framework for Recommendations at Pinterest</a></li> <li><a href="https://arxiv.org/abs/2205.11728">ItemSage: Learning Product Embeddings for Shopping Recommendations at Pinterest</a></li> <li><a href="https://arxiv.org/abs/2210.10634">RankT5: Fine-Tuning T5 for Text Ranking with Ranking Losses</a></li> <li><a href="https://arxiv.org/abs/2306.01599">Pretrained Language Model based Web Search Ranking: From Relevance to Satisfaction</a></li> <li><a href="https://arxiv.org/abs/2101.08751">Rethink Training of BERT Rerankers in Multi-Stage Retrieval Pipeline</a></li> <li><a href="https://arxiv.org/abs/2203.05765">Tevatron: An Efficient and Flexible Toolkit for Dense Retrieval</a></li> <li><a href="https://arxiv.org/abs/2207.02578">SimLM: Pre-training with Representation Bottleneck for Dense Passage Retrieval</a></li> <li><a href="https://arxiv.org/abs/2006.07733">Bootstrap your own latent: A new approach to self-supervised Learning</a></li> <li><a href="https://arxiv.org/abs/2002.07971">Gradient Boosting Neural Networks: GrowNet</a></li> <li><a href="https://arxiv.org/abs/2402.16829">GISTEmbed: Guided In-sample Selection of Training Negatives for Text Embedding Fine-tuning</a></li> <li><a href="https://arxiv.org/abs/2101.06983">Scaling Deep Contrastive Learning Batch Size under Memory Limited Setup</a></li> <li><a href="https://eugeneyan.com/writing/recsys-llm/?trk=feed_main-feed-card_feed-article-content">Improving Recommendation Systems & Search in the Age of LLMs</a></li> <li><a href="https://arxiv.org/abs/2311.11351">Scaling Law of Large Sequential Recommendation Models</a></li> <li><a href="https://arxiv.org/abs/2306.04833">Unified Embedding Based Personalized Retrieval in Etsy Search</a></li> <li><a href="https://arxiv.org/abs/2501.16450">360Brew: A Decoder-only Foundation Model for Personalized Ranking and Recommendation</a></li> <li><a href="https://arxiv.org/pdf/2409.02856">Building a Scalable, Effective, and Steerable Search and Ranking Platform</a></li> <li><a href="https://arxiv.org/abs/2501.12948">DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning</a></li> <li><a href="https://arxiv.org/abs/2412.19437">DeepSeek-V3 Technical Report</a></li> <li><a href="https://arxiv.org/abs/2407.21783">The Llama 3 Herd of Models</a></li> <li><a href="https://arxiv.org/abs/2501.16673">Adalflow</a></li> <li><a href="https://arxiv.org/abs/2506.15833">LLMs by dropping the text</a></li> <li><a href="https://arxiv.org/abs/2502.02013">Layer by layer</a></li> <li><a href="https://arxiv.org/abs/2505.24832">How much do LLMs memorize</a></li> <li><a href="https://arxiv.org/abs/2406.07496">TextGrad</a></li> <li><a href="https://arxiv.org/abs/2309.03409">LLMs as optimizers</a></li> <li><a href="https://arxiv.org/abs/2310.03714">DSPY</a></li> <li><a href="https://arxiv.org/abs/2211.01910">LLMs are human level prompt engineers</a></li> <li><a href="https://arxiv.org/abs/2210.11610">LLMs can self improve</a></li> <li><a href="https://arxiv.org/abs/2205.12548">RLPrompt</a></li> <li><a href="https://arxiv.org/abs/2207.00747">Rationale augmented Ensembles</a></li> <li><a href="https://arxiv.org/abs/2203.11171">Self Consistency Improves CoT</a></li> <li><a href="https://arxiv.org/abs/2201.11903">CoT Prompting Elicits Reasoning</a></li> <li><a href="https://arxiv.org/abs/2205.11916">LLMs are zero shot Reasoners</a></li> <li><a href="https://arxiv.org/abs/2205.10782">Instruction Induction</a></li> <li><a href="https://arxiv.org/abs/2009.06732">Efficient Transformers</a></li> <li><a href="https://arxiv.org/abs/2104.08786">Fantastically ordered prompts</a></li> <li><a href="https://arxiv.org/abs/2305.03495">Protegi</a></li> <li><a href="https://arxiv.org/abs/2210.03629">React</a></li> <li><a href="https://arxiv.org/abs/2311.16452">MedPrompt</a></li> <li><a href="https://arxiv.org/abs/1708.05123">DCNv1</a></li> <li><a href="https://arxiv.org/abs/2008.13535">DCNv2</a></li> <li><a href="https://arxiv.org/abs/2203.01941">RQ-VAE</a></li> <li><a href="https://arxiv.org/abs/2203.05556">Embedding Numeric Features</a></li> <li><a href="https://arxiv.org/abs/2411.07641">Top n-sigma</a></li> <li><a href="https://www.arxiv.org/abs/2508.09494">Learning facts at scale</a></li> <li><a href="https://arxiv.org/abs/2508.10478">Semantic IDs</a></li> <li><a href="https://arxiv.org/abs/2509.04259">RL's razor</a></li> <li><a href="https://arxiv.org/abs/1908.10396">Anisotropic Vector Quantization</a></li> <li><a href="https://arxiv.org/abs/2508.15884">Jet-Nemotron</a></li> <li><a href="https://www.arxiv.org/abs/2509.08827">Survey of RL</a></li> <li><a href="https://arxiv.org/abs/1711.00937">Discrete Neural Representation</a></li> <li><a href="https://arxiv.org/abs/2503.07891">Gemini Embedding</a></li> <li><a href="https://arxiv.org/abs/2508.21038">Theoretical Limits Embedding</a></li> <li><a href="https://arxiv.org/abs/2202.06991">Transformer memory</a></li> <li><a href="https://arxiv.org/abs/2408.10394">Unicorn</a></li> <li><a href="https://arxiv.org/abs/2406.10450">TokenRec</a></li> <li><a href="https://arxiv.org/abs/1807.05631">Joint Modelling of search and rec</a></li> <li><a href="https://arxiv.org/abs/2504.06714">Unifying search and rec</a></li> <li><a href="https://arxiv.org/abs/2311.09134">Scalable and Efficient Gen Info Retrieval</a></li> <li><a href="https://arxiv.org/abs/2010.00904">Autoregressive entity retrieval</a></li> <li><a href="https://arxiv.org/abs/2112.06905">MoE for language models</a></li> <li><a href="https://arxiv.org/abs/2101.03961">Switch Transformers</a></li> <li><a href="https://arxiv.org/abs/2012.14913">Feed forward are key value stores</a></li> <li><a href="https://arxiv.org/abs/1811.08008">End to end retrieval in continuous space</a></li> <li><a href="https://arxiv.org/abs/2105.02274">Rethinking search</a></li> <li><a href="https://arxiv.org/abs/2002.08910">How much knowledge</a></li> <li><a href="https://arxiv.org/abs/2108.08877">Sentence T5</a></li> <li><a href="https://arxiv.org/abs/2509.04501">GRAPE</a></li> <li><a href="https://arxiv.org/abs/2508.10925">GPT OSS</a></li> <li><a href="https://arxiv.org/abs/2509.12361">News Recommender</a></li> <li><a href="https://www.arxiv.org/abs/2510.02453">Advisor Models</a></li> <li><a href="https://arxiv.org/abs/2501.16673">LLM AutoDiff</a></li> <li><a href="https://arxiv.org/abs/2302.03668">Hard Prompts made Easy</a></li> <li><a href="https://arxiv.org/abs/2504.16929">I-Con</a></li> <li><a href="https://arxiv.org/abs/2510.13786">Scaling RL Compute</a></li> <li><a href="https://arxiv.org/abs/2509.17874">Nested Subspace Networks</a></li> <li><a href="https://arxiv.org/abs/2109.06141">Tilted Losses</a></li> <li><a href="https://arxiv.org/abs/2510.06477">Attention Sinks, Compression Valleys</a></li> <li><a href="https://arxiv.org/abs/2509.11208">Why LMs Hallucinate</a></li> <li><a href="https://arxiv.org/abs/2507.11768">LMs are Bayesian in Expectation, not Realization</a></li> <li><a href="https://arxiv.org/abs/2507.19457">GEPA</a></li> <li><a href="https://arxiv.org/abs/2111.02080">In Context Learning as Implicit Bayesian Inference</a></li> <li><a href="https://arxiv.org/abs/2107.05720">SPLADE</a></li> <li><a href="https://arxiv.org/abs/2410.06424">Rotation Trick</a></li> <li><a href="https://arxiv.org/abs/2107.03006">D3PM</a></li> <li><a href="https://arxiv.org/abs/2211.15029">DiffusionBERT</a></li> <li><a href="https://arxiv.org/abs/1611.01144">Gumbel-softmax</a></li> <li><a href="https://arxiv.org/abs/2206.10789">Parti Model</a></li> <li><a href="https://arxiv.org/abs/2202.04200">MaskGIT</a></li> <li><a href="https://arxiv.org/abs/2506.16698">SIDE</a></li> <li><a href="https://arxiv.org/abs/2411.18814">LIGER</a></li> <li><a href="https://arxiv.org/abs/2208.09257">Ultron</a></li> <li><a href="https://arxiv.org/abs/2305.02073">Understanding DSI</a></li> <li><a href="https://arxiv.org/abs/2103.03230">Barlow Twins</a></li> <li><a href="https://arxiv.org/abs/2510.04618">Agentic context engineering</a></li> <li><a href="https://arxiv.org/abs/2509.14252">LLM JEPA</a></li> <li><a href="https://arxiv.org/abs/2512.10942">VL JEPA</a></li> <li><a href="https://arxiv.org/abs/2506.14373">Discrete JEPA</a></li> <li><a href="https://arxiv.org/abs/2301.08243">I-JEPA</a></li> <li><a href="https://arxiv.org/abs/2505.17169">NTPS</a></li> <li><a href="https://arxiv.org/abs/2511.05963">Next latent prediction</a></li> </ul> <h1 id="rqvae"><a class="header" href="#rqvae">RQVAE</a></h1> <ul> <li><a href="https://arxiv.org/abs/2107.03312">Soundstream</a></li> <li><a href="https://arxiv.org/abs/1803.03382">Fast decoding with discrete latents</a></li> <li><a href="https://arxiv.org/abs/2309.15505">Finite Scalar Quantization</a></li> <li><a href="https://arxiv.org/abs/2401.14732">QINCo</a></li> <li><a href="https://arxiv.org/abs/2501.03078">QINCo2</a></li> <li><a href="https://arxiv.org/abs/2305.18378">Latent Quantize</a></li> <li><a href="https://arxiv.org/abs/2509.26469">DiVeQ</a></li> <li><a href="https://arxiv.org/abs/1906.00446">VQ-VAE2</a></li> </ul> <h1 id="semantic-ids-for-rec"><a class="header" href="#semantic-ids-for-rec">Semantic IDs for Rec</a></h1> <ul> <li><a href="https://arxiv.org/abs/2508.10584">DAS: Dual-Aligned Semantic IDs Empowered Industrial Recommender System</a></li> <li><a href="https://arxiv.org/abs/2410.09560">Scalable Semantic Representation</a></li> <li><a href="https://arxiv.org/abs/2406.10450">TokenRec</a></li> <li><a href="https://arxiv.org/abs/2304.04171">GenRet</a></li> <li><a href="https://arxiv.org/abs/2502.18965">OneRec</a></li> <li><a href="https://arxiv.org/abs/2506.13695">OneRec Technical Report</a></li> <li><a href="https://arxiv.org/abs/2203.13366">P5</a></li> <li><a href="https://arxiv.org/abs/2210.12316">VQRec</a></li> <li><a href="https://arxiv.org/abs/2305.15115">Semantic DSI</a></li> <li><a href="https://arxiv.org/abs/2404.08675">RecGPT</a></li> <li><a href="https://arxiv.org/abs/2502.18965">OneRec</a></li> <li><a href="https://arxiv.org/abs/2507.22224v1">SNAP GRID</a></li> <li><a href="https://arxiv.org/abs/2405.07314">Learnable item tokenization</a></li> <li><a href="https://arxiv.org/abs/2504.02137">Enhancing Embedding Representation Stability in Recommendation Systems with Semantic ID</a></li> <li><a href="https://arxiv.org/abs/2507.18800">Semantic IDs for Music Recommendation</a></li> <li><a href="https://arxiv.org/abs/2507.22224">Generative Recommendation with Semantic IDs: A Practitioner's Handbook</a></li> <li><a href="https://arxiv.org/abs/2310.19488">CoLLM: Integrating Collaborative Embeddings into Large Language Models for Recommendation</a></li> <li><a href="https://arxiv.org/abs/2311.09049">Adapting Large Language Models by Integrating Collaborative Semantics for Recommendation</a></li> <li><a href="https://arxiv.org/abs/2411.18814">Unifying Generative and Dense Retrieval for Sequential Recommendation</a></li> <li><a href="https://arxiv.org/abs/2305.06569">How to Index Item IDs for Recommendation Foundation Models</a></li> <li><a href="https://arxiv.org/abs/2008.07873">S^3-Rec: Self-Supervised Learning for Sequential Recommendation with Mutual Information Maximization</a></li> </ul> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="packages"><a class="header" href="#packages">Packages</a></h1> <h2 id="keybert-keyllm"><a class="header" href="#keybert-keyllm">KeyBERT, KeyLLM</a></h2> References: <ul> <li><a href="https://www.maartengrootendorst.com/blog/keybert/">KeyBERT Article</a></li> <li><a href="https://www.maartengrootendorst.com/blog/keyllm/">KeyLLM Article</a></li> </ul> KeyBERT and KeyLLM are packages to perform unsupervised keyword extraction from text. KeyBERT relies on BERT-based models, and the main idea is to extract n-gram phrases which have high semantic similarity to the overall document embedding. Some additional features are: <ul> <li>Allow user to specify phrase length for extraction</li> <li>Add diversification via MMR to get diverse phrases</li> </ul> KeyLLM taps on LLMs to enhance the keyword extraction. Basically, it creates a prompt to ask an LLM to extract keywords from a document. It integrates with KeyBERT such that we can use KeyBERT to cluster documents, and only run KeyLLM on one document per cluster to save costs. It can also use KeyBERT to suggest candidates and use the LLM to verify. <h2 id="pytorch-lightning"><a class="header" href="#pytorch-lightning">Pytorch Lightning</a></h2> The following pseudo code captures almost everything we need to know about pytorch lightning. Taken from <a href="https://lightning.ai/docs/pytorch/stable/common/lightning_module.html#hooks">here</a>. <pre><code class="language-python">def fit(self): configure_callbacks() if local_rank == 0: prepare_data() setup("fit") configure_model() configure_optimizers() on_fit_start() # the sanity check runs here on_train_start() for epoch in epochs: fit_loop() on_train_end() on_fit_end() teardown("fit") def fit_loop(): torch.set_grad_enabled(True) on_train_epoch_start() for batch in train_dataloader(): on_train_batch_start() on_before_batch_transfer() transfer_batch_to_device() on_after_batch_transfer() out = training_step() on_before_zero_grad() optimizer_zero_grad() on_before_backward() backward() on_after_backward() on_before_optimizer_step() configure_gradient_clipping() optimizer_step() on_train_batch_end(out, batch, batch_idx) if should_check_val: val_loop() on_train_epoch_end() def val_loop(): on_validation_model_eval() # calls `model.eval()` torch.set_grad_enabled(False) on_validation_start() on_validation_epoch_start() for batch_idx, batch in enumerate(val_dataloader()): on_validation_batch_start(batch, batch_idx) batch = on_before_batch_transfer(batch) batch = transfer_batch_to_device(batch) batch = on_after_batch_transfer(batch) out = validation_step(batch, batch_idx) on_validation_batch_end(out, batch, batch_idx) on_validation_epoch_end() on_validation_end() # set up for train on_validation_model_train() # calls `model.train()` torch.set_grad_enabled(True) </code></pre> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="skills"><a class="header" href="#skills">Skills</a></h1> This documents LinkedIn's approach to constructing a knowledge graph around skills. <h2 id="november-2022---building-linkedins-skills-graph-to-power-a-skills-first-world"><a class="header" href="#november-2022---building-linkedins-skills-graph-to-power-a-skills-first-world"><a href="https://www.linkedin.com/blog/engineering/skills-graph/building-linkedin-s-skills-graph-to-power-a-skills-first-world">November 2022 - Building LinkedIn's Skills Graph to Power a Skills-First World</a></a></h2> An important model for LinkedIn is the skills graph - it maps <code>39k skills</code> over <code>26 languages</code> and over <code>347k aliases</code>. The components of their structured skills graph: <ul> <li><code>39k</code> nodes, each of which is a skill</li> <li>Each skill has multiple aliases</li> <li>Skills are connected via edges which specify a hierarchical relationship. (e.g. <code>Marketing</code> is a <code>parent</code> of <code>Demand Generation</code> which is a parent of <code>Email Marketing</code>)</li> </ul> Skills Usage. Skills are automatically extracted from job postings and member histories. Job postings are allowed to tag up to 10 skills. Users can see something like <code>4/10 skills match your profile</code> to help them determine job fit. The skills graph is constructed both using ML and manual review by taxonomists. <h2 id="march-2023---building-and-maintaining-the-skills-taxonomy-that-powers-linkedins-skills-graph"><a class="header" href="#march-2023---building-and-maintaining-the-skills-taxonomy-that-powers-linkedins-skills-graph"><a href="https://www.linkedin.com/blog/engineering/data/building-maintaining-the-skills-taxonomy-that-powers-linkedins-skills-graph">March 2023 - Building and maintaining the skills taxonomy that powers LinkedIn's Skills Graph</a></a></h2> This article is more about how the Skills Graph is constructed. Each skill node includes the following foundational details (e.g. <code>Machine Learning</code>): <ul> <li>Description of the skill: <code>the study of computer algorithms...</code></li> <li>Aliases: <code>ML, ...</code></li> <li>Skills type: <code>Hard</code> or <code>Soft</code></li> <li>Skill ID</li> </ul> Each skill is represented by a <code>node</code> in the graph and nodes are linked via edges called <code>knowledge lineages</code>. Skills can relate for various reasons. Edges are directed and represent a parent-child relationship (e.g. <code>Software Development</code> -> <code>Back-end Software Development</code> -> <code>Python</code>) Each node can have multiple parents and/or children. <ul> <li>Multiple parents example: Both <code>Back-end Software Development</code> and <code>Mobile Software Development</code> are parents of <code>Java</code></li> <li>Multiple children example: <code>Supply Chain Management</code> has children <code>Supply Chain Engineering</code>, <code>Logistics Management</code> and <code>Digital Supply Chain</code></li> </ul> The hierarchical relationship allows us to enrich skills understanding, since if a person knows a particular skill, we may infer that he knows something about all the parent nodes and perhaps some of the "sibling" nodes. LinkedIn has certain quality guardrails on the structured skills, one of which is discouraging ambiguity. For example, an ambiguous skill like <code>Networking</code> may be mapped to skills like <code>Computer Networking</code> or <code>Professional Networking</code>. This type of ambigious relationship is not allowed, and LinkedIn either removes such edges or disallows such a node. The meaning of a phrase is determined by analyzing how the skill is used predominantly in LinkedIn. In cases where a phrase can have divergent meanings, the skill is disambiguated by expanding the phrase. For example, <code>Cadence</code> is disambiguated to <code>Cadence Software</code> and <code>Boundary</code> to <code>Boundary Line</code>. Architecture. The components to their ecosystem are as follows: <ol> <li>Human Curated KG. Presumably, this is a purely human-curated KG where all nodes and edges are taken as true, and constantly curated by human taxonomists.</li> <li>AI-generated KG. This is generated using ML models which use the human curated KG as training and validation data. The model behind this is KGBERT, which will be briefly covered below.</li> <li>Serving. Both the AI-generated KG and human-generated KG are made available to all LinkedIn services via a REST API for online serving and also on HDFS to power offline inference.</li> </ol> KGBERT is a model for automatic edge prediction, which can generate the AI-generated KG from the human curated one. The basic idea is that the human curated KG is used to generate training and validation data in the following form: <pre><code>[CLS] Tensorflow [SEP] Machine Learning [SEP] -> label: child - parent </code></pre> Two skills are concatenated to form the context. The skill is represented at random by its <code>title</code>, <code>description</code> or <code>title + description</code>. The context is fed into a BERT model and a linear + softmax layer is attached to the <code>[CLS]</code> token to generate a softmax probability over 3 options: <ul> <li>Parent -> Child</li> <li>Child -> Parent</li> <li>No relation</li> </ul> The positive labels are taken from edges in the human curated graph. The negative labels are generated by some heuristics: <ul> <li>Skill pairs from different industries</li> <li>Niece / Nephew pairs e.g. <code>Tensorflow</code> vs <code>Cognitive Computing</code></li> <li>Sibling pairs with the same parent e.g. <code>Tensorflow</code> vs <code>Pytorch</code></li> <li>Loosely related pairs that are 3 or more steps apart</li> </ul> The nice thing about this architecture is the clean separation of the human curated graph from the AI-generated one. The human-in-the-loop setup allows humans to review the AI-generated graph and add new edges to the human curated graph, which in turn improves the AI-generated graph. This seems superior over mixing both AI-generated edges and human-curated edges in the same graph. <h2 id="december-2023---extracting-skills-from-content-to-fuel-the-linkedin-skills-graph"><a class="header" href="#december-2023---extracting-skills-from-content-to-fuel-the-linkedin-skills-graph"><a href="https://www.linkedin.com/blog/engineering/skills-graph/extracting-skills-from-content">December 2023 - Extracting skills from content to fuel the LinkedIn Skills Graph</a></a></h2> This post takes a deeper look into how skills are extracted from data by LinkedIn from various contexts, such as job listings or member profiles. Note that skills can be mentioned either directly or indirectly (e.g. <code>you are expected to know how to apply different techniques to extract information from data and communicate insights through meaningful visualizations</code>). Hence, a simple span extraction approach will not be exhaustive in extracting skills. On the other hand, not constraining the problem to span extraction could lead to false positive errors, so our model has to be very accurate. <div class="table-wrapper"><table><thead><tr><th style="text-align: center"><img src="misc/../images/linkedin/skill_extraction_overview.png" alt="LinkedIn Skills Extraction" /></th></tr></thead><tbody> <tr><td style="text-align: center">An overview of the skills extraction architecture</td></tr> </tbody></table> </div> At a high level, the steps are: <ol> <li>The Skill Segmenter organizes the raw text into structured input, e.g. a resume is split into the <code>qualification</code>, <code>career history</code> etc. sections</li> <li>The next step aims to generate a candidate list of skills from the context: <ul> <li>2a. The Text Skill Tagger is a trie-based model that simply finds matches in the text that exactly match a node in the Skills Taxonomy. This puts the burden on the Skills Taxonomy to have all the aliases that cover every possible utterance of the skill. However, this model is very fast</li> <li>2b. The Semantic Tagger aims to overcome the coverage issues above and uses BERT semantic match to surface more candidates</li> <li>2c. The skills from 2a. and 2b. are expanded into more skills using the Skills Graph. It seems like they add immediate parents, children and siblings of each skill.</li> </ul> </li> <li>Each skill candidate is scored against the context to generate a confidence score. This section has a Shared Model and Domain-specific Model. <ul> <li>The Shared Model contains a <code>context-encoder</code>, which encodes the text surrounding the skill into an embedding, and a <code>entity-encoder</code>, which encodes surrounding skills, the job title, etc. All these embeddings are fed as context into the next stage. The Shared Model assumes that the relationship between the context surrounding each skill and the skill itself is constant across the different domains and thus benefit from a shared module. Their AB tests confirm this hypothesis and show that such multi-task training does increase lift in online metrics.</li> <li>Each domain (job posting, member profile, course) has its own scoring model. The embeddings from the Shared Model and each skill candidate are fed into the Domain-specific Model, which is presumably a form of cross encoder that generate a confidence score for each skill candidate. Finally, some threshold is applied and a final list of skills is generated.</li> </ul> </li> </ol> LinkedIn needs to generate skill extractions within <code>100ms</code> for a high volume of edits. Hence, they used knowledge distillation to distil the BERT model down 80% in parameters. As we can see, the LinkedIn architecture is fairly complex and context specific. A natural question for smaller players is how we can tap on LLMs today to simplify parts of this architecture to achieve comparable performance. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="hash-collisions"><a class="header" href="#hash-collisions">Hash Collisions</a></h1> Bloom embeddings is one way of handling large number of user / item IDs. Instead of assigning each unique ID to a unique embedding, we may perform hashing to assign each unique ID to one or multiple bins. The embeddings at these positions are then taken and typically summed to get the final representation of each user / item. One natural question arises: given <code>n</code> number of unique values and <code>m</code> number of bins, what is the expected number of bins to have 2 or more values assigned to it? This collision is especially a problem if we only represent each unique value with one embedding (i.e. <code>num_hashes = 1</code>). In practice, we would usually have at least <code>2</code> or more hashes to avoid this problem. We may approach the problem by observing that each item is hashed uniformly and has a m1 chance of landing in a particular bin. The hashing is also independent, i.e. it is not affected by the hashing for any other item. This means that the expected number of items assigned to a particular bin may be modeled as a X∼Binomial(n,m1) distribution, since we have <code>n</code> trials and a fixed probability of "success" of <code>1/m</code>. <ul> <li>The probability we desire is P(X≥2)=1−P(X=0)−P(X=1)</li> <li>The PMF is P(X=k)=(kn)(m1)k(1−m1)n−k, so: <ul> <li>P(X=0)=(1−m1)n</li> <li>P(X=1)=mn⋅(1−m1)n−1</li> </ul> </li> </ul> So we may put the above together to obtain P(X≥2). It then remains to get the expected number of colliding bins by m⋅P(X≥2), by using the linearity of expectation. For example, given n=50,000,m=1,000,000, the expected number of colliding bins is <code>1,210</code> (out of <code>1 million</code> bins), which is not insignificant. Now note that since n is large and p=1/m is small, X may be well approximated by Y∼Poisson(n/m). Recall that the Poisson PMF is P(Y=k)=λk⋅e−λ/k!. So: <ul> <li>P(Y=0)=e−n/m</li> <li>P(Y=1)=mn⋅e−n/m</li> </ul> Using this approximation, we get the expected number of colliding bins as <code>1,209</code>, which is a very good approximation. Hence the formula we want is: Expected bins with collision=m×(1−e−n/m−mn⋅e−n/m) <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="imo-2025-q6"><a class="header" href="#imo-2025-q6">IMO 2025 Q6</a></h1> <blockquote> Q6: Consider a 2025×2025 grid of unit squares. Matilda wishes to place on the grid some rectangular tiles, possibly of different sizes, such that each side of every tile lies on a grid line and every unit square is covered by at most one tile. Determine the minimum number of tiles Matilda needs to place so that each row and each column of the grid has exactly one unit square that is not covered by any tile. </blockquote> This is an interesting problem with a surprising solution. It's also an easy problem to understand and think about in your head (though certainly not easy to solve!). I have an interactive visualization at <a href="https://charleslow.github.io/grid_visualization">https://charleslow.github.io/grid_visualization</a> to explore this problem. <h2 id="initial-attempt"><a class="header" href="#initial-attempt">Initial Attempt</a></h2> The first instinct is to take an approach like below. That is, we observe that using single row / column rectangular tiles is sufficient to tile the grid. Hence the solution uses 2(n−1) tiles if the grid is n×n. It's n−1 instead of n because of the gaps (black squares) up against the grid boundaries. I thought this was the solution initially, and tried to prove it but failed. <figure style="max-width:320px;margin:1.25rem auto;padding:0 1.5rem;text-align:center;"> <img src="misc/../images/imo_2025_q6/naive_example.png" alt="Naive 4-tile configuration" style="width:100%;height:auto;border:1px solid #ccc;border-radius:6px;box-shadow:0 2px 4px rgba(0,0,0,0.08);" /> <figcaption style="font-size:0.85rem;color:#555;margin-top:0.4rem;">Figure: A naive approach.</figcaption> </figure> <h2 id="my-solution"><a class="header" href="#my-solution">My Solution</a></h2> Let's draw a vertical and horizontal line from each uncovered square (call it a "gap"). Observe there may be <code>2-4</code> such lines emanating in the NSEW directions from each gap. Overall, there would be 4n−4 such lines (−4 because there is no line at the first & last row / column). Let's call each of these lines a "half-line". Since each <code>row, col</code> pair has one gap, we are assured that a tiling that covers all the half-lines satisfies Matilda (as it will then cover every row and column, except for the gaps). Now the question is, how quickly can we cover the 4n−4 half-lines with a minimal number of tiles? In other words, how many half-lines can we clear with each tile we introduce? Now, it is tricky to analyze the number of lines we clear on each turn, because it depends on the preceding steps. For example, if we had cleared all but one remaining tile on each half-line (show example), then the next step can clear a large number of half-lines. It is hard to reason across multiple turns. So an important observation: <blockquote> Lemma 1: The number of gaps that each tile touches is an upper bound for the number of half-lines cleared after k steps. Proof. This is quite self-explanatory, as touching a gap is pre-requisite to clearing a half-line that emanates from that side of the gap. </blockquote> Lemma 1 shifts the analysis from clearing half-lines to analyzing (and maximizing) the number of gaps we touch with every tile we introduce. The number of gaps touched per tile is easier to reason about, as we shall see. We will make the upper bound an equality later on with a specific strategy. <blockquote> Lemma 2: The number of gaps touched per tile is capped at <code>4</code>. Proof. This is also not hard to see. Imagine a tile bounded on all sides by gaps. Since there can only be one gap per row and column, the tile cannot be touching more than one gap per side. Since the tile only has 4 sides, it can only touch up to 4 gaps. </blockquote> <h3 id="4-x-4-configuration"><a class="header" href="#4-x-4-configuration">4 x 4 Configuration</a></h3> To get a minimal tiling, we wish to maximize the number of "4-tiles" (i.e. tiles touching <code>4</code> holes). This may be done by minimizing the size for a 4-tile at 2×2. One such configuration is shown below. <figure style="max-width:320px;margin:1.25rem auto;padding:0 1.5rem;text-align:center;"> <img src="misc/../images/imo_2025_q6/optimal_4_tile.png" alt="Optimal 4-tile configuration" style="width:100%;height:auto;border:1px solid #ccc;border-radius:6px;box-shadow:0 2px 4px rgba(0,0,0,0.08);" /> <figcaption style="font-size:0.85rem;color:#555;margin-top:0.4rem;">Figure: A minimal 2×2 tile touching four gaps.</figcaption> </figure> A naive solution is to just replicate this pattern repeatedly along the diagonal, like so. Since we take <code>5</code> tiles to tile each 4×4 mini-grid, this gives us 45n+2(4n−1) tiles for a n×n grid. For e.g. for a 12×12 grid shown, the naive solution gives us 2(12−1)=22, but our new strategy gives 5×3+2(3−1)=19. So we already have a lower number than the original naive solution, but can we do better? <figure style="max-width:320px;margin:1.25rem auto;padding:0 1.5rem;text-align:center;"> <img src="misc/../images/imo_2025_q6/repeated_optimal_4_tile.png" alt="Optimal 4-tile configuration" style="width:100%;height:auto;border:1px solid #ccc;border-radius:6px;box-shadow:0 2px 4px rgba(0,0,0,0.08);" /> <figcaption style="font-size:0.85rem;color:#555;margin-top:0.4rem;">Figure: Repeat across diagonal.</figcaption> </figure> <h3 id="powers-of-4"><a class="header" href="#powers-of-4">Powers of 4</a></h3> Staring at the figure above, we can see that putting all our 4×4 mini-grids along the diagonal is to commit the same mistake as our original naive solution. More specifically, if we treat each 4×4 mini-grid as a single square, we can see that the problem reduces once again to the original problem, except that the number of rows and columns shrink by a factor of <code>4</code>. Therefore, we can squeeze more <code>4-tiles</code> out by using the same strategy (except that we are now treating each 4×4 mini-grid as an atomic square). An example using 16×16 is shown: <figure style="max-width:320px;margin:1.25rem auto;padding:0 1.5rem;text-align:center;"> <img src="misc/../images/imo_2025_q6/optimal_16_by_16.png" alt="Optimal 16x16 configuration" style="width:100%;height:auto;border:1px solid #ccc;border-radius:6px;box-shadow:0 2px 4px rgba(0,0,0,0.08);" /> <figcaption style="font-size:0.85rem;color:#555;margin-top:0.4rem;">Figure: Optimal 16 x 16 tiling.</figcaption> </figure> We can recursively use this strategy for grid sizes that are powers of 4. The minimal number of tiles for a grid of 4i×4i would be given by this recursive formula (for i≥1, and T0=0): Ti=4×Ti−1+5 Which after expanding gives us the following formula for a grid of 4k×4k: Tk=5i=0∑k−14i For example, a grid of 4×4 requires 5 tiles, a grid of 42×42 requires 5×(4+1)=25 tiles, etc. <h3 id="2025-is-not-a-power-of-4"><a class="header" href="#2025-is-not-a-power-of-4">2025 is not a power of 4</a></h3> What if our grid is not a power of 4? We can split it up into a series of grids that are powers of 4. So we create a simple algorithm to do this. Suppose we start with a grid of size n0×n0. <ul> <li>At initialization, set our current grid size n=n0</li> <li>Find the largest n′ such that n′=4k≤n, and partition out the n′×n′ sub-grid</li> <li>Find the minimal number of tiles required for this sub-grid and add it to our current count</li> <li>Add 2 tiles to extend the sides to cover the full n×n grid</li> <li>Set n=n−n′ and continue until n<4</li> <li>Use the naive solution for n<4</li> </ul> Applying this algorithm to our original problem of 2025×2025, we have: 2025=45+3(44)+3(43)+2(42)+2(41)+40 So using the formula (and letting T′=T+2 to account for the additional 2 tiles for each sub-grid) we have: T2025=T5′+3T4′+3T3′+2T2′+2T1′+T0′=3,379 Turns out this solution is incorrect. The assumption that the minimum 2×2 tiling to get a 4-tile turns out to be suboptimal. <h3 id="appendix-code-to-compute-number-of-tiles"><a class="header" href="#appendix-code-to-compute-number-of-tiles">Appendix: Code to compute number of tiles</a></h3> <pre><code class="language-python">import math def tiles(k: int): return int(5 * sum(math.pow(4, i) for i in range(k))) def largest_power(n: int): assert n >= 1, "Input must be a positive integer." k = 0 while 4 ** k <= n: k += 1 return k - 1 def main(n0: int): count = 0 n = n0 while n > 0: k = largest_power(n) count += tiles(k) + 2 n -= 4 ** k assert n == 0, f"Remaining value is {n} != 0." return count print(main(2025)) </code></pre> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="identities"><a class="header" href="#identities">Identities</a></h1> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="sigmoid"><a class="header" href="#sigmoid">Sigmoid</a></h1> <h2 id="sigmoid-relationship-to-bradley-terry-model"><a class="header" href="#sigmoid-relationship-to-bradley-terry-model">Sigmoid Relationship to Bradley-Terry Model</a></h2> Suppose we have scores si,sj for objects i and j. Under the Bradley-Terry model, we can express the preference probability as follows: P(i▹j)=esi+esjesi=(esi+esj)/esi1=1+esj/esi1=1+esj−si1=1+e−(si−sj)1=σ(si−sj) Where σ(x)=1/(1+e−x) is the sigmoid function. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="statistics-1"><a class="header" href="#statistics-1">Statistics</a></h1> <h2 id="adding--subtracting-rvs"><a class="header" href="#adding--subtracting-rvs">Adding / Subtracting RVs</a></h2> Let X,Y be two independent random variables. Then Var(X±Y)=Var(X)+Var(Y). First, show that Var(X)=E(x2)−μX2 where x∼X. Proof: Var(X)=E[] Reference: https://apcentral.collegeboard.org/courses/ap-statistics/classroom-resources/why-variances-add-and-why-it-matters <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> Summaries of individual papers that I have taken a deeper look into. <h1 id="quick-reads"><a class="header" href="#quick-reads">Quick Reads</a></h1> Here I put down quick summaries of papers I read quickly. May possible revisit these with its own page if a deep dive is warranted. <h2 id="jepa"><a class="header" href="#jepa">JEPA</a></h2> <h3 id="assran-2023---self-supervised-learning-from-images-with-a-joint-embedding-predictive-architecture"><a class="header" href="#assran-2023---self-supervised-learning-from-images-with-a-joint-embedding-predictive-architecture"><a href="https://arxiv.org/abs/2301.08243">Assran 2023 - Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture</a></a></h3> This is the I-JEPA paper. This paper aims to perform JEPA-style self-supervised learning for vision transformers without using image augmentations (like rotation, shear etc.). The main idea is to take a random image patch as context and aim to predict multiple other image patches as target. Importantly: <ul> <li>The loss is to minimize l2-distance between the Predictor(Context Encoding) and Target Encoding in latent space</li> <li>To avoid representation collapse in latent space, the target encoder is updated as a trailing exponential moving average of the Context Encoder</li> <li>This approach is simple and scalable and matches performance of SOTA methods using image augmentations</li> </ul> <h3 id="geng-2022---recommendation-as-language-processing-rlp-a-unified-pretrain-personalized-prompt--predict-paradigm-p5"><a class="header" href="#geng-2022---recommendation-as-language-processing-rlp-a-unified-pretrain-personalized-prompt--predict-paradigm-p5"><a href="https://arxiv.org/abs/2203.13366">Geng 2022 - Recommendation as Language Processing (RLP): A Unified Pretrain, Personalized Prompt & Predict Paradigm (P5)</a></a></h3> This is one of the earliest papers to reframe recommendation task as a language modelling task. The idea is to represent user-item interaction data and other metadata in recommendation tasks as an <code>input</code> to <code>response</code> task for supervised fine-tuning (or pre-training). Then we can have a unified language model to perform various tasks in a recommendation system. The tasks proposed are: <ul> <li>Sequential Recommendation: Given purchase history <code>item_123, item_456</code>, what is the next item? <code>1581</code></li> <li>Rating Prediction: What star rating would <code>user_123</code> give to <code>item_456</code>? 5.0</li> <li>Explanation Generation: What explanation or review would <code>user_123</code> give to <code>item_456</code> (iphone case)? You can use it to protect your phone</li> <li>Review Summarization: Give a short sentence summarizing this review <code><insert review></code>. <code><insert summary></code></li> <li>Direct recommedation: Pick the most suitable item from the following list to recommend to <code>user_123</code>: <code>item_123, item_456, ...</code>. <code>item123</code></li> </ul> These tasks are mixed together and then used to pretrain an encoder-decoder model from scratch, and then used for various recommendation tasks. <h3 id="tang-2023---se-dsi"><a class="header" href="#tang-2023---se-dsi"><a href="https://arxiv.org/abs/2305.15115">Tang 2023 - SE-DSI</a></a></h3> This is an expansion of the <a href="./papers/tay_2022.html">Differentiable Search Index</a> paper. The argument is that teaching an LLM to learn arbitrary document ids is not semantic. Hence we should represent documents as semantic words to aid learning. Specifically: <ul> <li>Docid Representation. For each document, use a T5 model to generate a query to serve as document identifier (5% collision).</li> <li>Input features. Use 3 types of input features to predict the "query docid": <ul> <li>Whole document</li> <li>Key passages identified using textrank</li> <li>Key sentenes identified using textrank</li> </ul> </li> </ul> This approach beats vanilla DSI, as we use semantic information as document representation. Note that evaluation task is Q&A. <h3 id="chen-2023---understanding-dsi-for-text-retrieval"><a class="header" href="#chen-2023---understanding-dsi-for-text-retrieval"><a href="https://arxiv.org/abs/2305.02073">Chen 2023 - Understanding DSI for Text Retrieval</a></a></h3> Analyzes the shortcomings of DSI through 3 metrics: <ul> <li>Exclusivity: There should be a one-to-one mapping between document contents and semantic ID. <ul> <li><code>Task</code>: Given first N tokens of document, can we retrieve the correct ID</li> </ul> </li> <li>Completeness: The LLM should remember as much information about each document as possible. <ul> <li><code>Task</code>: Use BERT model to find important chunks in document, then use these chunks as query to see if we can retrieve the correct ID.</li> </ul> </li> <li>Relevance: Can the LLM rank documents in relevance order to a query. <ul> <li><code>Task</code>: For a given query, check the document at position <code>k</code> of the LLM's recommendation order to see how relevant it is. Specifically, measure the cross encoder score of document dk against the query and compare to score of a random document dr. Lower score is considered a failure</li> <li>Found that DSI only generates relevant documents at position <code>1</code>, and then becomes highly irrelevant thereafter.</li> </ul> </li> </ul> Proposed improvements: <ul> <li>Data Filtering: Break each document into chunks and use a teacher model <code>TCT-Colbert</code> to perform retrieval using that chunk as query. Only successful retrievals that get the document back are considered useful content used for LLM fine-tuning.</li> <li>Multi-task learning: Normal DSI indexing task only seeks to predict one document ID given its contents or a query. They propose to add an additional task where we get LLM to learn to predict a list of relevant document IDs. <ul> <li>The list of relevant document IDs are generated using a teacher model (<code>TCT-Colbert</code> again).</li> </ul> </li> </ul> <h3 id="ju-2025---grid"><a class="header" href="#ju-2025---grid"><a href="https://arxiv.org/abs/2507.22224">Ju 2025 - GRID</a></a></h3> Snap's GRID repository which is a modular repo for semantic ID based generative retrieval to facilitate architecture comparison. A generative retrieval system has the following components: <ul> <li>Semantic ID Tokenization: <ul> <li>Convert item description to embedding</li> <li>Quantize embedding to tokens (K-means, VQ-VAE, RQ-VAE)</li> </ul> </li> <li>Sequential Recommender: <ul> <li>Whether to use user token</li> <li>Data augmentation using sliding windows</li> <li>Architecture: Encoder-decoder vs decoder only</li> <li>Beam search</li> </ul> </li> </ul> Note: this system only considers the case where each user is represented as a sequence of items, and the LLM is trained from scratch (not finetuned from pretrained). Their experiments are interesting: <ul> <li>Residual K-means tokenization is best, compared to VQ-VAE and RQ-VAE. This is quite surprising and I'm not 100% convinced.</li> <li>Larger OTS embedding models are better, no surprise.</li> <li>Encoder-decoder sequential recommender is clearly better. Somewhat surprising, I guess there is some inductive bias that is useful in this setting to have full attention over the user sequence.</li> <li>Sliding window is better. Sliding window means that given a training data sequence of a user represented by item interactions, we generate all possible contiguous sequences to augment our dataset. Not surprising.</li> <li>SID de-duplication is unnecessary. Between appending an arbitrary token to deduplicate SIDs or just randomly selecting from SIDs with multiple items hashed to it, the performance difference is negligible. This is assuring that we don't need to further complicate the system.</li> <li>Constrained beam search is unnecessary. Constrained beam search means that we maintain a trie to only generate tokens leading to a valid SID. The results show that this has negligible performance difference to unconstrained beam search, which has much better efficiency. This is assuring that we can just use uncontrained beam search.</li> </ul> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="weinberger-2009---hashing-for-multitask-learning"><a class="header" href="#weinberger-2009---hashing-for-multitask-learning">Weinberger 2009 - Hashing for Multitask Learning</a></h1> <a href="https://arxiv.org/abs/0902.2206">Weinberger 2009 - Feature Hashing for Large Scale Multitask Learning </a> This paper proposes a method to represent a large feature set in a smaller space by using hashing. It shows analytically that with a sufficiently large hash dimension m: <ul> <li>The inner product between instances is preserved, i.e. doing a dot product between instances in the hashed dimension approximates the true dot product in the original dimension</li> <li>The same applies to learning a weight vector to generate predictions in the hashed space: the error of approximation goes to zero as M increases</li> </ul> <h2 id="setup-1"><a class="header" href="#setup-1">Setup</a></h2> Consider having data points x(1),...,x(n)∈Rd, where d can be very large (e.g. millions). This setting is easily realized when we use, for example, word bi-grams and tri-grams as term-frequency features to perform some kind of text classification task. Such a large feature vector is unwieldy, and also inefficient since the feature vector is very sparse for a given text. The hashing trick maps the high dimensional input vector to a smaller dimension feature space with the notation ϕ:X→Rm, such that m<<d. We start with the following definitions: <ul> <li>Let h be a hash function h:N→{1,...,m}</li> <li>Let E be a hash function E:N→{±1}</li> </ul> Note that while the definitions map from an input integer, we may apply them to texts as well, since any finite-length text may be assigned to a unique integer. This is typically done in practice by applying some hash algorithm to a given string, and then using the modulo function to restrict it to the desired range. With this, and given two vectors x,x′∈Rd, we define the hash feature map: ϕj(h,E)(x)=i∈Z : h(i)=j, 1≤i≤d∑E(i)xi Where j∈1,...,m is an index in the hashed dimension space, and i∈1,...,d is an index in the input dimension space. We get a hash collision if more than one i term is hashed into a given position j. For brevity, we may just write ϕj(h,E):=ϕ. <h2 id="analysis"><a class="header" href="#analysis">Analysis</a></h2> With this setup, the paper aims to prove analytically that hashing in this manner preserves the characteristics of the original space. In other words, we can significantly reduce the dimension of our features but achieve the same predictive effect as the original space by doing the hashing trick. This also means that the detrimental effect of hash collisions is minimal with a sufficiently large m. We won't trace through all the results, just the important and simple ones. Lemma 2 The hash kernel is unbiased, i.e. Eϕ[⟨x,x′⟩ϕ]=⟨x,x′⟩. Proof. The proof simply starts by expanding the inner product in the hashed space as follows: ⟨x,x′⟩ϕ=i=1,...,d∑j=1,...,d∑E(i)E(j)⋅xixj′⋅δh(i),h(j) Where δh(i)=h(j) is an indicator variable which takes 1 if h(i)=h(j) (i.e. they are hashed to the same position) and 0 otherwise. To see that this expansion is true, consider a position in the hashed space, e.g. k. The value at position k looks something like the following. We just need to move the summands to the left and use the δ variable to denote the common hash positions where xi and xj′ interact (if <code>i</code> and <code>j</code> are hashed to different positions, they clearly do not interact in an inner product). [⟨x,x′⟩ϕ]k=<svg xmlns="http://www.w3.org/2000/svg" width='0.667em' height='3.600em' viewBox='0 0 667 3600'><path d='M403 1759 V84 H666 V0 H319 V1759 v0 v1759 h347 v-84 H403z M403 1759 V0 H319 V1759 v0 v1759 h84z'/></svg>i∈Z : h(i)=k, 1≤i≤d∑E(i)xi<svg xmlns="http://www.w3.org/2000/svg" width='0.667em' height='3.600em' viewBox='0 0 667 3600'><path d='M347 1759 V0 H0 V84 H263 V1759 v0 v1759 H0 v84 H347z M347 1759 V0 H263 V1759 v0 v1759 h84z'/></svg><svg xmlns="http://www.w3.org/2000/svg" width='0.667em' height='3.600em' viewBox='0 0 667 3600'><path d='M403 1759 V84 H666 V0 H319 V1759 v0 v1759 h347 v-84 H403z M403 1759 V0 H319 V1759 v0 v1759 h84z'/></svg>j∈Z : h(j)=k, 1≤j≤d∑E(j)xj<svg xmlns="http://www.w3.org/2000/svg" width='0.667em' height='3.600em' viewBox='0 0 667 3600'><path d='M347 1759 V0 H0 V84 H263 V1759 v0 v1759 H0 v84 H347z M347 1759 V0 H263 V1759 v0 v1759 h84z'/></svg> Now note that we can decompose the expectation over ϕ into its independent constituents, i.e. h and E respectively (since the two hashes are independent): Eϕ[⟨x,x′⟩ϕ]=Eh[EE⟨x,x′⟩ϕ] Now we just need to observe that the hashed values E(i),E(j) are independent from all other terms in general, but also independent from each other whenever i=j (provided our hash function is pairwise independent). Thus when i=j, the summand is: EE[E(i)]⋅EE[E(j)]⋅EE[xixj′⋅δh(i),h(j)] These are clearly 0 because E[E(i)]=0. So the original summation reduces to: Eϕ[⟨x,x′⟩ϕ]=Eϕ<svg xmlns="http://www.w3.org/2000/svg" width='0.667em' height='3.600em' viewBox='0 0 667 3600'><path d='M403 1759 V84 H666 V0 H319 V1759 v0 v1759 h347 v-84 H403z M403 1759 V0 H319 V1759 v0 v1759 h84z'/></svg>i=1,...,d∑E(i)2⋅xixi′<svg xmlns="http://www.w3.org/2000/svg" width='0.667em' height='3.600em' viewBox='0 0 667 3600'><path d='M347 1759 V0 H0 V84 H263 V1759 v0 v1759 H0 v84 H347z M347 1759 V0 H263 V1759 v0 v1759 h84z'/></svg>=⟨x,x′⟩ Not only is the hashed inner product unbiased, it also has a variance that scales down in O(m1). The proof does a similar but more tedious expansion as the above, and assumes that x,x′ have l2-norm of 1. This suggests that the hashed inner product will be concentrated within O(m<svg xmlns="http://www.w3.org/2000/svg" width='400em' height='1.08em' viewBox='0 0 400000 1080' preserveAspectRatio='xMinYMin slice'><path d='M95,702 c-2.7,0,-7.17,-2.7,-13.5,-8c-5.8,-5.3,-9.5,-10,-9.5,-14 c0,-2,0.3,-3.3,1,-4c1.3,-2.7,23.83,-20.7,67.5,-54 c44.2,-33.3,65.8,-50.3,66.5,-51c1.3,-1.3,3,-2,5,-2c4.7,0,8.7,3.3,12,10 s173,378,173,378c0.7,0,35.3,-71,104,-213c68.7,-142,137.5,-285,206.5,-429 c69,-144,104.5,-217.7,106.5,-221 l0 -0 c5.3,-9.3,12,-14,20,-14 H400000v40H845.2724 s-225.272,467,-225.272,467s-235,486,-235,486c-2.7,4.7,-9,7,-19,7 c-6,0,-10,-1,-12,-3s-194,-422,-194,-422s-65,47,-65,47z M834 80h400000v40h-400000z'/></svg>1) of the true value. These results are sufficient to justify use of the hashed inner product space in practice. That is, we can perform recommendations in the hashed space with sufficiently large m (we can tune that using validation error) to make the large feature space tractable. The paper goes on to prove more detailed bounds on the error and norm which are of less practical significance. <h2 id="multi-task-learning"><a class="header" href="#multi-task-learning">Multi-task Learning</a></h2> The authors argue that this method is especially useful in the multi-task learning setting. Consider an email spam classification task where the vocab space is V and the user space is U. The parameter space is thus V×U, i.e. we wish to learn a user-specific weight vector wu∈R∣V∣ for each user u, which allows us to personalize the spam filter for each user (different users have slightly differing definitions of what is spam). The authors suggest the following approach: <ul> <li>Use the hashing trick to hash each term v into the hashed space. e.g. <code>data</code> is passed into a global hash function ϕ0 and assigned to a position</li> <li>Each user gets his/her own hash function ϕu. This may be implemented by using the same hash function but appending the <code>user_id</code> like so: <code>user1_data</code>, which hashes the same term into a new position.</li> <li>We may thus represent each instance by ϕ0(x)+ϕu(x)∈Rm, capturing both a global element (some terms are universally spam-indicative) and a personalized element (some terms are specifically indicative for a user)</li> <li>Finally, we learn a weight parameter wh∈Rm by training it in the hashed space</li> </ul> Empirically, for their task of ∣V∣=40 million, ∣U∣=400,000, they found that performance starts to saturate with m≈4 million. This is a very small fraction of the total space ∣V∣×∣U∣, showing the effectiveness of their method. Nevertheless, we should note that <code>4 million</code> is still a rather large space. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="rendle-2009---bayesian-personalized-ranking"><a class="header" href="#rendle-2009---bayesian-personalized-ranking">Rendle 2009 - Bayesian Personalized Ranking</a></h1> <a href="https://arxiv.org/abs/1205.2618">BPR: Bayesian Personalized Ranking from Implicit Feedback</a> This is one of the most cited papers for collaborative filtering. It proposes a pairwise learning algorithm for item recommendation from implicit feedback that remains one of the most performant to date. It is pitched as a rival method to Weighted Matrix Factorization, the other strong model of the day. BPR argues that it is superior to WMF because it explicitly optimizes for ranking. <h2 id="setup-2"><a class="header" href="#setup-2">Setup</a></h2> Let U, I denote the set of all users and items respectively. We have access to an implicit feedback data S⊂U×I. The task is to produce each user with a personalized ranking >u⊂I2. Note that I2 denotes the cartesian product of I with itself, so it represents the set of all ordered pairs of elements of I. The ordering >u is a subset of these pairs where a preference relationship is indicated by the model. For convenience we also denote: <ul> <li>Iu+:={i∈I:(u,i)∈S}</li> <li>Ui+:={u∈U:(u,i)∈S}</li> </ul> In implicit feedback systems, only positive interactions between user and item and observed. The typical approach is to put the interactions in a matrix and fill in unobserved entries with 0. A model is then fitted to this model. The paper makes an interesting observation that this problem is ill-defined, since a perfectly expressive model that fits the training data perfectly will fail to make any prediction at all, since all unobserved entries will be given a score of 0. Hence regularization methods are employed to avoid such overfitting. The idea of BPR is to avoid making judgment on the pairwise preference between two items with the same score. That is, if items i and j are both interacted by user u, we cannot judge if one is preferred over the other. Also, if i and j are both unobserved interactions, we cannot make such judgement either. Thus the training data DS:I×I×I is denoted by the following. DS:={(u,i,j)∣i∈Iu+ AND j∈I∖Iu+} Note that this definition means that (u,i,j)∈DS⟹(u,j,i)∈/DS, since j cannot be positive for u if it was included as a negative. <h2 id="bpr-loss"><a class="header" href="#bpr-loss">BPR Loss</a></h2> The bayesian formulation for finding a ranking is to find the model parameters Θ that maximize the following probability: p(Θ∣>u)∝p(>u∣Θ)⋅p(Θ) We assume that: <ul> <li>All users act independently of each other</li> <li>The ordering of each pair of items (i,j) for a given user is independent of the ordering of every other pair</li> </ul> Then, we can write the likelihood across all users as: u∈U∏p(>u∣Θ)=(u,i,j)∈U×I×I∏p(i>uj∣Θ)δ((u,i,j)∈DS)⋅[1−p(i>uj∣Θ)]δ((u,j,i)∈DS) Where δ is the indicator function for the preference relationship. In other words, the likelihood of the overall ordering is the product of the likelihood of each u,i,j triplet. For each u,i,j triplet, the likelihood is given by the model's prediction for given label 0 or 1. Note that the above term is only not equals to 1 if i>uj or j>ui. Also, since we observed above that (u,i,j)∈DS⟹(u,j,i)∈/DS and vice versa, only one of the two terms will come into play for each entry. The above formula can be simplified to: u∈U∏p(>u∣Θ)=(u,i,j)∈DS∏p(i>uj∣Θ) Note: Not too sure about the above step, would have thought that the 1−p term should also come into play when j>ui. <h2 id="model"><a class="header" href="#model">Model</a></h2> Now we can model the preference probability using a model as: p(i>uj∣Θ):=σ(y^uij(Θ)) Where y^ denotes a BPR-agnostic model that generates a predicted real-valued score for the triplet (u,i,j). Note that σ(x):=1/(1+exp(−x)) is the sigmoid function. To complete the bayesian modelling, a prior distribution p(Θ) is introduced over the model parameters. For simplicity, we let p(Θ) take the form of a multivariate normal distribution with zero-mean and all covariances zero, with equal variance for each parameter (i.e. λΘ∈R): p(Θ)∼N(0,λΘI) The maximum posterior estimator for Θ is thus given by the following: argmaxΘlnp(>u∣Θ)⋅p(Θ)=ln(u,i,j)∈DS∏σ(y^uij)⋅p(Θ)=(u,i,j)∈DS∑lnσ(y^uij)+lnp(Θ)=(u,i,j)∈DS∑lnσ(y^uij)+λΘ∣∣Θ∣∣2 Note that in the last step, the log prior is translated into L2 regularization. This is apparently well studied and I will explore the derivation at a later time. A possible resource is <a href="https://medium.com/intuition/understanding-l1-and-l2-regularization-with-analytical-and-probabilistic-views-8386285210fc#66b3">Yuki Shizuya - Understanding L1 and L2 regularization with analytical and probabilistic views</a>. The paper suggests using SGD updates (i.e. batch size of <code>1</code>) by randomly sampling triplets with replacement from DS. This converges much faster than doing user-wise SGD. Finally, we've assumed y^ is model-agnostic up to now. For matrix factorization, the authors suggest the following form. In other words, we compute the prediction for each user, item pair (usually via dot product of embeddings, but can be anything) and take the difference. y^uij=y^ui−y^uj <h2 id="cornac-implementation"><a class="header" href="#cornac-implementation">Cornac Implementation</a></h2> Cornac has a <a href="https://github.com/PreferredAI/cornac/blob/master/cornac/models/bpr/recom_bpr.pyx">Cython implementation of BPR</a> that is fast (but not memory scalable when number of user and items is large). <code>num_threads</code> can be increased for faster parallel processing (Cython overrides the GIL). At initialization: <ul> <li>User embedding <code>U</code> of shape <code>n_users, k</code> is drawn randomly from a standard uniform distribution where <code>k</code> is the embedding dimension</li> <li>Then, we take (U−0.5)/k. I suppose this leads to small, centered values which leads to stable initial training. Normalizing by <code>k</code> ensures that the dot product between vectors does not explode.</li> <li>Item embedding of shape <code>n_items, k</code> is initialized similarly</li> <li>Optional bias vector of shape <code>n_items</code> is initialized to the zero vector</li> </ul> Note that <code>train_set.X</code> is a <code>scipy.sparse.csr_matrix</code>, which has 3 main vectors: <ul> <li><code>data</code>: values of non-zero entries</li> <li><code>indices</code>: column indices of the non-zero values</li> <li><code>indptr</code>: index pointers to the start of each row in <code>data</code> and <code>indices</code></li> </ul> The vector <code>user_ids</code> is constructed from the <code>csr_matrix</code>. It is of the same length as <code>X.indices</code> and represents the row indices of the non-zero values. Thus we have both row and column indices of the non-zero values. The main training loop samples a <code>(u, i, j)</code> triplet randomly each turn. It does so by the following steps: <ul> <li>A random index <code>i</code> between <code>0</code> and <code>len(user_ids)</code> is generated</li> <li>The <code>user_id</code> and <code>item_i_id</code> are obtained by indexing <code>user_ids[i]</code> and <code>X.indices[i]</code> respectively</li> <li>Now a random index <code>j</code> between <code>0</code> and <code>n_items</code> is generated</li> <li>The <code>item_j_id</code> is obtained by indexing <code>neg_item_ids[j]</code> where <code>neg_item_ids = np.arange(n_items)</code></li> <li>A check is performed on the sparse matrix that <code>item_j_id</code> is not a positive item for <code>user_id</code>. If so, we skip this triplet.</li> <li>Pointers to the start of the relevant <code>u, i, j</code> embeddings are then obtained</li> <li>Note that <code>1</code> epoch comprises <code>len(user_ids)</code> number of triplets</li> </ul> The Cython code then computes and manually applies the SGD updates derived in the paper. Will omit since we do not need to manually compute the updates if using autograd. Note: This implementation of BPR assumes binary relevance (an item is interacted or not). It does not allow for finer-grained preference relationships (e.g. a floating point rating score to rank items), which in theory BPR does support. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="from-ranknet--lambdarank--lambdamart"><a class="header" href="#from-ranknet--lambdarank--lambdamart">From RankNET → LambdaRank → LambdaMART</a></h1> <a href="https://www.microsoft.com/en-us/research/uploads/prod/2016/02/MSR-TR-2010-82.pdf">Paper Link</a>. This is an overview paper that explains the model behind LambdaMART, a technique used to learn Gradient Boosted Trees that optimize for NDCG on recommendation-type of datasets. <h2 id="ranknet"><a class="header" href="#ranknet">RankNET</a></h2> RankNET is a generic framework for training a learn to rank model. Given a differentiable model f:Rd↦R that produces a score from an n-dimensional input vector, RankNET is able to train f such that for a given query session with items i=1,...n and corresponding features xi∈Rd:i=1,...,n, f learns to predict the relationship f(xi)>f(xj) for any two items i,j in the query session when i is a better recommendation than j. The differentiable model f is typically a neural network or boosted trees. RankNET uses revealed pair-wise preferences within each query session to train f. Specifically, suppose we have a data for one query session as follows: <div class="table-wrapper"><table><thead><tr><th>query</th><th>item</th><th>clicked</th></tr></thead><tbody> <tr><td>qid1</td><td>a</td><td>0</td></tr> <tr><td>qid1</td><td>b</td><td>1</td></tr> <tr><td>qid1</td><td>c</td><td>0</td></tr> </tbody></table> </div> We can transform this data into a pairwise dataset as follows, where yjk denotes the preference relationship between itemj and itemk which we inferred from the click data. Note that the pairwise comparisons are only made within the same query session (e.g. <code>qid1</code>), as it reflects a given user's preferences in the context of a query and the items impressed to him/her in that session. <div class="table-wrapper"><table><thead><tr><th>query</th><th>itemj</th><th>itemk</th><th>yjk</th></tr></thead><tbody> <tr><td>qid1</td><td>a</td><td>b</td><td>-1</td></tr> <tr><td>qid1</td><td>a</td><td>c</td><td>0</td></tr> <tr><td>qid1</td><td>b</td><td>a</td><td>1</td></tr> <tr><td>qid1</td><td>b</td><td>c</td><td>0</td></tr> </tbody></table> </div> The pairwise setting is now more amenable to modelling (compared to directly optimizing for a good ranking), since we can now treat the task as a classification problem. For each row of the pairwise dataset, we only need to model the probability that itemj is preferred (or not) to itemk. This can be formalized using a cross entropy loss comparing the predicted preference of our model to the revealed preference in the dataset. First, we model the predicted probability from the model. Given row i of the pairwise dataset and itemj and itemk respectively, we model the predicted probability that itemj is preferred to itemk (using ▹ to denote a preference relationship) by passing the score difference between the predicted scores y^ij:=f(xij) and y^ik:=f(xik) for items j and k respectively through a sigmoid function, like so: P^ijk:=P^(itemij ▹ itemik)=1+exp(−a(y^ij−y^ik))1 Now let us denote the revealed probability that itemj is preferred to itemk as Pijk such that: <ul> <li>Pijk=1 if we prefer item j to item k</li> <li>Pijk=0.5 if we have no preference between the two items</li> <li>Pijk=0 if we prefer item k to item j</li> </ul> The cross entropy loss of our model can then be expressed as: L=i∑[−PijklogP^ijk−(1−Pijk)log(1−P^ijk)] For convenience, let us denote yijk:=2Pijk−1 (and conversely, Pijk=2yijk+1), which translates into the following: <ul> <li>yijk=1 if we prefer item j to item k</li> <li>yijk=0 if we have no preference between the two items</li> <li>yijk=−1 if we prefer item k to item j</li> </ul> Let us also define the convenience variable z:=−logP^ijk=log[1+exp(−a(y^ij−y^ik))]. The cross entropy loss then simplifies to: L=i∑[−PijklogP^ijk−(1−Pijk)log(1−P^ijk)]=i∑[21(1+yijk)z−21(1−yijk)[(−a)(y^ij−y^ik)−z]]=i∑[z+2a(1−yijk)(y^ij−y^ik)] Note that in line 2 of the above, we use the useful identity that log(1−sigmoid(x))=x+log(sigmoid(x)) and logP^ijk=−z. In line 3, the first and last term of line 2 cancel out to simply return z. Having written out the loss function, we now need to differentiate the loss with respect to the model scores and parameters to obtain the gradient descent formula used to train the RankNET model. Differentiating L wrt y^ij and y^ik gives: ∂y^ij∂L=2a(1−yijk)+1+exp(−a(y^ij−y^ik))−a⋅exp(−a(y^ij−y^ik))=2a(1−yijk)−1+exp(a(y^ij−y^ik))a=−∂y^ik∂L Note that the first line of the above uses the result dxdlnf(x)=f(x)f′(x). We obtain line 2 by multiplying the right term by exp(a(y^ij−y^ik)) in both the numerator and denominator. We obtain line 3 by observing that L is a function of d:=y^ij−y^ik, such that ∂y^ij∂L=∂d∂L⋅∂y^ij∂d=∂d∂L and likewise ∂y^ik∂L=∂d∂L⋅∂y^ik∂d=−∂d∂L. The symmetry of the derivative wrt j and k will be important for the next section on factorizing RankNET to speed it up. Finally, we use the gradient to update individual parameters wl∈R of the model f. In the below, n denotes the number of data points in the pairwise dataset. This update procedure rounds up the discussion on RankNET and is sufficient for training a generic differentiable model f from ranking data. wl←wl−η∂wl∂L=wl−nηi∑[∂y^ij∂L⋅∂wl∂y^ij+∂y^ik∂L⋅∂wl∂y^ik] <h2 id="factorizing-ranknet"><a class="header" href="#factorizing-ranknet">Factorizing RankNET</a></h2> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="schroff-2015---facenet"><a class="header" href="#schroff-2015---facenet">Schroff 2015 - FaceNET</a></h1> <a href="https://ar5iv.labs.arxiv.org/html/1503.03832">Schroff 2015 - FaceNet: A Unified Embedding for Face Recognition and Clustering</a> This paper proposes the Triplet Loss for learning a face recognition system. The task is that given a large number of person identities with a number of images associated with each person, to learn a model representation for an image in euclidean space, such that images belonging to the same person are close by and images for different persons are far away. This paper was impactful because it improved the SOTA on face verification by a large margin. At this point, representation learning often trained a CNN classification model to classify images to a known identity. A bottleneck layer of relatively low dimensionality in the middle of the network is chosen as the representation of an image. In contrast to this indirect method, this paper directly optimizes the representation using contrastive learning. <h2 id="setup-3"><a class="header" href="#setup-3">Setup</a></h2> Let the embedding of image x be represented by f(x)∈Rd, and constrain ∣∣f(x)∣∣2=1. Now for a given person <code>i</code>, we want the anchor image xia to be closer to other images of the same person xip (positive) than images of any other person xin (negative), with some margin α. That is, we desire a model f such that: ∣∣f(xia)−f(xip)∣∣22+α<∣∣f(xia)−f(xin)∣∣22 ∀(xia,xip,xin)∈T Where T denotes the set of all possible triplets in the dataset. The authors note that most triplets would easily satisfy the desired constraint, hence we need some way to picking good triplets to optimize the learning process. The loss we desire to minimize is thus as below. The intuition is that we want the anchor-positive distance dap to be small but the anchor-negative distance dan to be small. L=i=1∑N[∣∣f(xia)−f(xip)∣∣22 − ∣∣f(xia)−f(xin)∣∣22+α]+ With this loss function, we see that we have 3 types of triplets: <ul> <li>Easy triplets where α<dan−dap. These will contribute 0 loss since L=dap−dan+α<0 will be a negative value, and the positive operator makes the loss 0. This makes sense because the model already categorizes these triplets well and there is nothing more to learn from these examples.</li> <li>Semi-hard triplets where 0<dan−dap≤α. These will contribute a small loss, since 0≤L=dap−dan+α<α. These are triplets where the positive is already closer to the anchor than the negative (which is good), but the distance is still within the margin so we want to make the distance greater. In other words, .</li> <li>Hard triplets where dan−dap<0. These will contribute a large loss. These are triplets where the negative is closer to the anchor than the negative.</li> </ul> <h2 id="triplet-selection"><a class="header" href="#triplet-selection">Triplet Selection</a></h2> As with most contrastive learning methods, the authors observe that the selection of triplets is important for learning well. In an ideal world, we would want to select the hardest negatives for each anchor, i.e.: xin=argminxin ∣∣f(xia)−f(xin)∣∣2 However, this is undesirable because: <ul> <li>We may end up mostly selecting mislabelled instances or corrupted images</li> <li>It is computationally infeasible to search across the whole corpus</li> <li>Trying to learn hard negatives at the start of training may cause learning to collapse</li> </ul> Hence, the authors propose: <ul> <li> Mini batch negatives. Instead of finding the hardest negatives across the entire corpus, they mine for the hardest negatives in the mini-batch. This improves computational efficiency and mitigates the mislabelled/corrupted data issue. </li> <li> Curriculum Learning. At the start of training, the authors select easier triplets, and gradually increase the difficulty as training goes along. Presumably, this is done by thresholding negatives based on dan−dap>t and starting from high t to low t. </li> <li> Semi-hard negatives. The authors mention that it is helpful to limit negatives to dan−dap>0, but it is unclear whether they always do this or only at the start of training. The intuition is to cap the difficulty of the triplets, as triplets that are too difficult may actually hinder learning due to corrupted data or mislabelling. Note: In a later paper <a href="https://openaccess.thecvf.com/content_CVPR_2020/html/Lee_Large_Scale_Video_Representation_Learning_via_Relational_Graph_Clustering_CVPR_2020_paper.html">Lee 2020 - Large Scale Video Representation Learning via Relational Graph Clustering</a>, it implies that FaceNET uses semi-hard negatives throughout the entire training process, which makes sense given that the paper warns against using hard negatives. In light of this, I would presume FaceNET searches for the hardest semi-hard negatives in each mini batch as negatives for each triplet. </li> </ul> <h2 id="application-to-semantic-search"><a class="header" href="#application-to-semantic-search">Application to Semantic Search</a></h2> The triplet loss may be applied directly to semantic search. However, it is important to note that the paper assumes that for each person, all others-labelled instances are negatives. This is a suitable assumption for face recognition as each image can only belong to one person, but it is not true for semantic search, where a document may be relevant for multiple queries. Hence, the mislabelling issue when mining hard negatives is amplified for semantic search. I imagine that the selection of accurate negatives for semantic search would require more verification and filtering. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="covington-2016---deep-nns-for-youtube-recs"><a class="header" href="#covington-2016---deep-nns-for-youtube-recs">Covington 2016 - Deep NNs for Youtube Recs</a></h1> <a href="https://static.googleusercontent.com/media/research.google.com/en//pubs/archive/45530.pdf">Deep Neural Networks for YouTube Recommendations</a> This is an old paper that documents YouTube's recommender system, but quite foundational in the Recsys world. This paper also occurs at a time when there was a migration of models into deep learning methods, and Tensorflow was open sourced in the year before. <h2 id="retrieval"><a class="header" href="#retrieval">Retrieval</a></h2> As with most recommender systems, there is a candidate generation and ranking step. The candidate generation surfaces "hundreds" of videos. The predecessor to this work is a matrix factorization approach trained under rank loss (assuming it is <a href="papers/./rendle_2009.html">BPR</a>). This work may be thought of as generalizing the MF approach to non-linear transformations of arbitrary features. The paper views recommendation as an extreme multiclass classification problem, where the aim is to predict the probability that the video watch wt at time t is video i. In the below, V is the corpus of videos, U is a specific user and C is a specific context. Also u,v∈RN are high dimensional embeddings representing user and item respectively. P(wt=i∣U,C)=∑j∈VevjTueviTu It is well known that the above is intractable with millions of classes, so this paper samples negative classes from the background distribution and then corrects for the sampling via <a href="papers/./yi_2019.html">importance weighting</a>. This paper samples "several thousands of negatives" for each true label, and then minimizes the softmax loss for the true label against the negatives. Note: it is not specified what is the background distribution, it could either be a uniform distribution over items or the empirical marginalized impression probability of each item over a fixed time period. Also note that the modern approach is to use in-batch negatives for efficiency, as opposed to the purely random negatives approach here. At serving time, a fast ANN search algorithm is used to find the top scoring items. This paper used a hashing algorithm in <a href="https://proceedings.mlr.press/v28/weston13.html">Weston 2013 - Label Partitioning for Sublinear Ranking</a> for this, which is probably an outdated approach today. <h3 id="retrieval-features"><a class="header" href="#retrieval-features">Retrieval: Features</a></h3> Each user is represented by features such as: <ul> <li>Watch history, represented by a variable-length sequence of video IDs. Each video is represented by an embedding, and then the element-wise average is taken such that the whole history is represented by a dense embedding. Note: this is prior to the attention mechanism, the modern approach is to run the sequence through self attention before taking the element-wise average to get a more nuanced representation.</li> <li>Search history, represented by a variable-length sequence of search terms, treated similarly to watch history. Note that in the experiments, both watch and search history were capped at <code>50</code> terms each.</li> <li>Categorical features such as Geographic location are embedded into a dense vector.</li> <li>Simple binary and float features such as gender, log in status are input directly into the network as real values normalized to [0, 1].</li> </ul> The features are all concatenated together to form a long dense vector representing the raw user inputs. e.g. if watch and search history are length n, concatenating would give a vector of 2n and so on. The dense vector is passed into a fully connected feed forward network with shrinking layers to get a final fixed size dense vector u representing the user. Similar approach is done for items. <div class="table-wrapper"><table><thead><tr><th style="text-align: center"><img src="papers/../images/covington_2016_retrieval_architecture.png" alt="Candidate Retrieval Architecture" /></th></tr></thead><tbody> <tr><td style="text-align: center">Candidate Retrieval Architecture</td></tr> </tbody></table> </div> <h3 id="retrieval-example-age-feature"><a class="header" href="#retrieval-example-age-feature">Retrieval: Example Age Feature</a></h3> Since video popularity typically spikes in the first few days after upload and dies down after, an important feature is the "example age" to capture the "freshness" of an item. Since we typically have a training window of several weeks, failing to account for the age of an item would mean that the model will learn to predict the average watch probability of a video over the entire training period. This is not ideal since it makes the model biased toward older videos (which have had more time for exposure) than newer videos. The authors correct for this by including the <code>age of the training example</code> as a feature during training. Suppose our training window is <code>90 days</code>, then this feature would range from <code>90</code> (or <code>1</code> after normalization) for training examples on the first day of the training window to <code>0</code> for training examples on the last day of the training window. At serving time, the value of this feature is set to <code>0</code> to reflect the fact that we are making predictions at the very end of the window. The authors show that this approach helps the model learn to predict the time-sensitive effect of popularity for a given item that matches the distribution in the data. Note: This is different from incorporating the <code>age of the item</code> as a feature during training. I suppose while <code>age of the item</code> can capture the same information, it requires more work at serving time, since we need to set the feature value to the actual <code>age of the item</code> at the time of serving. In contrast, the above approach allows us to just set the feature to <code>0</code> which is more elegant. <h3 id="retrieval-surrogate-problem"><a class="header" href="#retrieval-surrogate-problem">Retrieval: Surrogate Problem</a></h3> The paper emphasizes that recommendation involves solving a surrogate problem that is then transferred to live inference. Hence care must be taken to ensure that we do not overfit the surrogate problem which can hurt live metrics. Some insights are: <ul> <li>Generate training examples from all YouTube video watches, not just from YouTube's recommendations, so that relevant videos can be propagated quickly even if they did not originate on YouTube itself.</li> <li>Generate a fixed number of training examples per user, so that very active users do not dominate the loss function. Not sure how this works in practice, presumably a fixed number of examples per user (or just 1?) is sampled for each training epoch. This significantly improved live metrics.</li> <li>Withhold information from the classifier. This is a counter-intuitive point that shows us how tricky it is to find a good surrogate problem. Given the structure of YouTube, users can find videos either from homepage recommendations or from searching for a video. In the latter case, if a user searches for say <code>taylor swift</code>, it is typically followed by user watches of <code>taylor swift</code> videos. However, we do not want to reproduce this behaviour for homepage recommendations, such that a user is always shown search results based on his / her last search query - this would perform very poorly! Hence the representation of search terms using a simple element-wise average (instead of using a sequence model like RNNs) actually leads to better live results.</li> <li>Predict the next watched video. It is important that training is not performed on random held-out samples, but that the next video is always held out and a "rollback" of features available prior to the video watch is supplied to the classifier. This is because video watching is highly sequential in nature, e.g. episodic videos are watched in sequence and users explore artists starting with the popular videos before moving to niche ones, hence it is important that the surrogate problem captures this behaviour.</li> </ul> <h3 id="retrieval-ablation-studies"><a class="header" href="#retrieval-ablation-studies">Retrieval: Ablation Studies</a></h3> The paper explored a few settings. At <code>depth 0</code>, the MLP simply transforms the dense input vector in one step to the output dimension of <code>256</code>, i.e. essentially a linear approach. For <code>depth 1</code>, an in-between ReLU layer of <code>256</code> is introduced; for <code>depth 2</code> in-between ReLU layers of <code>512 -> 256</code> is introduced and so on. The performance increases significantly from <code>depth 0</code> to <code>depth 1</code> with diminishing benefits as we increase the model size. In terms of features, it is interesting that incorporating searches is very important for performance. <ul> <li><code>Watches only</code> has MAP of around 6%</li> <li><code>Watches + Searches</code> has MAP of around 9%</li> <li><code>Watches + Searches + Example Age</code> has MAP of around 10%</li> <li><code>All features</code> has MAP of around 11%</li> </ul> <h2 id="ranking"><a class="header" href="#ranking">Ranking</a></h2> The difference between ranking and retrieval is that ranking aims to specialize predictions to the specific setting in which impressions are shown. In practice, this means that the ranker has access to many more variables about the video and the user's relationship to the video than the retriever. The general architecture for the ranking model is similar to the retrieval model, with the difference that the impressed video is also included as part of the features and the model simply outputs a logit for prediction. <div class="table-wrapper"><table><thead><tr><th style="text-align: center"><img src="papers/../images/covington_2016_ranking_architecture.png" alt="Ranking Architecture" /></th></tr></thead><tbody> <tr><td style="text-align: center">Ranking Architecture</td></tr> </tbody></table> </div> <h3 id="ranking-feature-engineering"><a class="header" href="#ranking-feature-engineering">Ranking: Feature Engineering</a></h3> Despite the promise of deep learning, this paper found that significant time was still spent engineering features due to the challenge of encoding features for recommender systems. The main challenge for the authors was in representing a sequence of temporal user interactions and how these actions relate to the video being impressed. Note: The modern approach seems to move away from such intense feature engineering. Instead, the past interactions and video being impressed are simply embedded and passed into a self attention mechanism. The model can then learn interaction features in the hidden layers. The downside of the modern approach is paying higher compute cost at inference time, but it may still be worthwhile given the significant engineering effort required for generating and retrieving hundreds of features manually. The paper highlighted User-item previous interaction features as the most important and impactful. This includes features such as: <ul> <li> How many videos from this channel did the user watch before? </li> <li> When was the last time the user watched a video on this topic? These continuous features describing past user interactions on related items are especially powerful </li> <li> Features describing number of past impressions of this (user, item) pair are important for a responsive recommender system. For example, if the user was shown this item many times recently but did not interact with it, the recommender can learn to demote this item so that the user's recommendations can be "refreshed". Note: For YouTube, they maintain such impression features at up-to-the-second latency because of its importance. For other use cases, a latency of say 15 minutes might still be quite helpful. </li> <li> Feature scores from candidate generation step are also useful </li> </ul> Similar to the retrieval step, categorical features are represented by dense embeddings, with a separate embedding table per categorical feature. The vocabulary for each feature is determined by a single pass over the training data at training time, and the vocabulary size is capped at <code>1 million</code> based on impression count. Out of vocabulary values during inference time are simply mapped to the zero embedding. Multivalent features (e.g. a sequence of video IDs for watch history) are averaged element-wise before being concatenated. The paper also found that normalizing continuous features was critical for performance. Specifically, a feature value x with empirical distribution f will be normalized to x~=∫−∞xdf, which is the cumulative probability of the empirical distribution up to value x. To make it tractable, the authors compute the quantiles of the empirical distribution and do linear interpolation between the quantiles. For e.g. if the empirical distribution of x is <code>(1, 2, 8, 10)</code>, then <code>1</code> will be mapped to <code>0.25</code>, <code>2</code> to <code>0.5</code> and so on. Note: this is different from min max normalization, which only uses the min and max value. <h3 id="ranking-predicting-watch-time"><a class="header" href="#ranking-predicting-watch-time">Ranking: Predicting Watch Time</a></h3> The authors found that predicting watch time instead of a binary watch variable gets much better performance at live metrics. This is incorporated into the model in the form of weighted logistic regression. For positive examples (where the user did watch the video), the example is weighted by the amount of time the user spent watching the video. For negative examples (where the user did not watch), the example gets unit weight. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="recommendations-as-treatments"><a class="header" href="#recommendations-as-treatments">Recommendations as Treatments</a></h1> <a href="https://arxiv.org/pdf/1602.05352.pdf">Paper Link</a>. Training and evaluation data for recommender systems is subject to selection bias, because the probability of observing a data point depends on (i) users interacting with items which they like and (ii) the recommender system pushing items which they think the user likes. This leads to data Missing Not at Random (MNAR) and leads to biased model training and evaluation. Consider users interacting with movies. Denote users u∈{1,...,U} and movies i∈{1,...,I}. Let Y∈RU×I denote the true rating / interaction matrix between user and item, and Y^ the predicted matrix. Let O∈{0,1}U×I be the observability matrix of whether an interaction was observed, and P∈RU×I be the probability of observation, i.e. Pu,i=P(Ou,i=1). Ideally, we want the propensity matrix <code>P</code> to be uniform across all entries, which gives us the Missing Completely at Random (MCAR) condition. <h2 id="evaluating-only-on-observed-entries-is-biased"><a class="header" href="#evaluating-only-on-observed-entries-is-biased">Evaluating only on observed entries is biased</a></h2> Given a model Y^, we often want to compute evaluation metrics δu,i(Y^,Y) on how well Y^ approximates Y. For example: δu,i(Y^,Y)=(Yu,i−Y^u,i)2 MSE δu,i(Y^,Y)=logrank(Y^u,i)Yu,i DCG The Risk function may then be denoted: R(Y^)=U×I∑u∑iδu,i(Y^,Y) However, R(Y^) cannot be computed because most entries in Y are missing. Typically, we estimate R^ using only observed entries. This estimator is naive because it simply assumes the propensity matrix P is uniform. This assumption is often false in recommender systems because a small set of popular items tend to get impressed a lot. Hence this estimator will favour models that lean towards the popular items compared to models that recommend rarer / new items. R^naive(Y^)=∣{(u,i):Ou,i=1}∣1⋅u,i:Ou,i=1∑δu,i(Y^,Y) <h2 id="unbiased-estimators"><a class="header" href="#unbiased-estimators">Unbiased estimators</a></h2> The main idea in this paper is to view this problem as analogous to estimating treatment effects in causal inference. Think of recommending an item as an intervention analogous to treating a patient with a specific drug. The goal is to estimate the outcome of a new recommendation (clicked or not) or new treatment (recovered or not), while most outcomes between <code>u,i</code> pairs are not known. The key to resolving this bias is to understand the assignment mechanism that generated O, namely the propensity matrix P. We can then correct for the propensity. We need to assume that Pu,i>0∀u,i, i.e. full support, because otherwise the IPS estimator below is undefined. <h2 id="ips-estimator"><a class="header" href="#ips-estimator">IPS Estimator</a></h2> The main estimator is the Inverse Propensity Score (IPS) estimator. Assuming that the assignment mechanism P is known: R^ips(Y^∣P)=∣{(u,i):Ou,i=1}∣1⋅u,i:Ou,i=1∑Pu,iδu,i(Y^,Y) We have simply normalized each score by its inverse propensity, so that popular items with a high chance of being shown get their score reduced proportionally (and likewise in the opposite direction for rare items). We can show that the IPS estimator is unbiased as we take the expectation over the random observability matrix. That is, suppose we are allowed to sample an infinitely large number of observations based on Pu,i, the average of the IPS estimator over these datasets will be the true risk function. This simply happens because the expected value of Ou,i is simply the propensity, and if we know the propensity we can just cancel it out. EO[R^ips(Y^∣P)]=U⋅I1u∑i∑EOu,i[Pu,iδ_u,i(Y^,Y)⋅Ou,i]=U⋅I1u∑i∑δu,i(Y^,Y)=R(Y^) <h2 id="comments"><a class="header" href="#comments">Comments</a></h2> <ul> <li>This paper only addresses exposure bias, but not position bias. In other words, it assumes that all impressions are equal, which is valid when a user is shown an ad, but not when a user is shown a list of search results. In the latter case, items in lower positions have a lower probability of being considered by the user, even though both are considered to be observed.</li> </ul> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="doersch-2016---tutorial-on-vaes"><a class="header" href="#doersch-2016---tutorial-on-vaes">Doersch 2016 - Tutorial on VAEs</a></h1> <a href="https://arxiv.org/abs/1606.05908">Tutorial on Variational Autoencoders</a> Tutorial on VAEs from a computer vision perspective. <h2 id="introduction"><a class="header" href="#introduction">Introduction</a></h2> Generative modelling deals with models of distributions P(X), defined over data points X∈X. For example, X may represent pixels of an image. P(X) may indicate that a set of pixels that look like a real image has high probability, whereas images looking like random noise get low probability. For generative modelling to be useful, we don't just want an unconditional distribution. Instead, we want to generate images from a conditional distribution, e.g. images that look like another image or following a certain caption. <h2 id="latent-variable-models"><a class="header" href="#latent-variable-models">Latent Variable Models</a></h2> Generative modelling usually starts with a latent variable. For example, in digit image generation, we may first want to choose a particular digit (say <code>5</code>) to generate. The image model then knows to generate pixels corresponding to this digit (as opposed to deciding on the fly). z is typically chosen to indicate the latent variable (which can be multi-dimensional). Formally, we have a vector of latent variables z∈Z which we can easily sample from based on some probability density function P(z) defined over Z. In VAEs, this distribution is typically a multi-variate standard gaussian distribution. We now want to make sure that we have some function f that maps from the latent variable space Z to the output space X. Say we have a family of deterministic functions (typically some neural network) f(z;θ) which is parametrized by a vector θ∈Θ, where f:Z×Θ→X. In this setup, randomness is injected by the random variable z such that f(z;θ) is a random variable in the space X. To actually generate images that look like X, we wish to optimize θ such that we can sample z from P(z) and with high probability, f(z;θ) will look like that X's in our data. How do we formalize this notion? Naively, f(z;θ) just gives a random point prediction. We need some probability distribution P(X∣z;θ) which tells us how likely a given training sample X is under the generative model f. By the law of total probability, we can then write down the Maximum Likelihood objective: P(X)=∫P(X∣z;θ)P(z)dz The choice of P(X∣z;θ) for VAEs is often gaussian. Specifically, the most standard choice is: P(X∣z;θ)=N(X∣f(z;θ),σ2×I) This means that the function f(z;θ) gives us the gaussian mean, and the covariance matrix is a fixed diagonal (as σ is a hyperparameter). Notice that this gaussian expression allows us to express the idea that we just need f(z;θ) to produce samples that look like X, but does not have to exactly match some X. This smoothness is critical for generative modelling — if we specify P(X∣z;θ) to be the dirac delta function (i.e. all the probability mass is on the specific output produced by f(z;θ)), it would be impossible to learn from data, as the function is zero almost everywhere. Also note that while gaussian is the most common choice, it does not have to be so. We just need P(X∣z) to be computable and be continuous in θ. For example, if X is binary, then P(X∣z) can be a bernoulli parametrized by f(z;θ) (although it's unfathomable why one would do this). <h2 id="variational-autoencoders"><a class="header" href="#variational-autoencoders">Variational Autoencoders</a></h2> In order to maximize P(X) above, there are two big problems to solve: <ul> <li>How do we define the latent variables z and what information they encode</li> <li>How to deal with the integral over z</li> </ul> <h3 id="problem-1-how-to-represent-latents"><a class="header" href="#problem-1-how-to-represent-latents">Problem 1: How to represent Latents?</a></h3> For the first problem, VAEs basically prescribe minimal structure to the latents z, and say that the samples of z can be drawn from a simple distribution, say N(0,I). Then, the onus falls on the model f to map from a simple gaussian distribution to the complex distribution which describes our data. We know empirically that this is not a problem for an arbitrarily large neural network. <h3 id="problem-2-how-to-deal-with-integral-over-pz"><a class="header" href="#problem-2-how-to-deal-with-integral-over-pz">Problem 2: How to deal with integral over P(z)?</a></h3> For the second problem, the naive approach is to deal with the integral via sampling. That is we sample a large number of latents z={z1,...,zn} and compute P(X)≈n1∑iP(X∣zi). The problem with this approach is that in high dimensional spaces, we need a very large n to get an accurate estimate of P(X). This is because for most instances of z, P(X∣z) will be very close to 0. The key idea behind VAEs is to speed up sampling by attempting to sample values of z that are likely to have produced X, and compute P(X) just from those. To do this, we need a new function Q(z∣X) which when given an image X, produces a distribution over z values likely to have produced X. If the space of z values that are likely under Q is smaller than the space of z values that are likely under P(z), we will be able to estimate Ez∼QP(X∣z) much more cheaply. However, by introducing a new arbitrary distribution Q, we are no longer sampling under P(z), so we cannot directly obtain P(X) from it. Thus we need some way to relate Ez∼QP(X∣z) and P(X). This relationship is one of the cornerstones of variational Bayesian methods. Let us start by defining KL divergence between P(z∣X) and some arbitrary distribution Q(z) (which may or may not depend on X). D[Q(z)∣∣P(z∣X)]=Ez∼Q[logQ(z)−logP(z∣X)] Recall that KL divergence is asymmetric, and measures how different two probability distributions are. In this case, the expectation is taken over the distribution of z values under Q. Now apply Bayes rule to P(z∣X), and do some rearranging: D[Q(z)∣∣P(z∣X)]logP(X)−D[Q(z)∣∣P(z∣X)]=Ez∼Q[logQ(z)−logP(X∣z)−logP(z)]+logP(X)=Ez∼Q[logP(X∣z)]−D[Q(z)∣∣P(z)] Some comments on the above: <ul> <li>P(X) comes out as a constant because it does not depend on z.</li> <li>We can group Q(z) and P(z) together into its own KL divergence term.</li> </ul> So far, we have not made any assumption on the arbitrary distribution Q(z). In the context of trying to maximize P(X), it makes sense to construct a Q which does depend on X. So we make that dependency on X explicit. Let's call this the ELBO equation. logP(X)−D[Q(z∣X)∣∣P(z∣X)]=Ez∼Q[logP(X∣z)]−D[Q(z∣X)∣∣P(z)] This equation is core to the VAE, so we should understand it deeply. <ul> <li>The left hand side represents what we want to optimize: <ul> <li>logP(X) was the original maximum likelihood objective - we want our model to produce images that look like X</li> <li>D[Q(z∣X)∣∣P(z∣X)] is the error or deviation of our tractable, estimated distribution Q(z∣X) from the true, intractable oracle distribution P(z∣X). This term is always more than or equals to 0 and is 0 if and only if Q=P.</li> </ul> </li> <li>The right hand side is called the Evidence Lower Bound (ELBO). In bayesian statistics, the marginal P(X) is called the evidence, because our data X is evidence for how good our model is. The RHS is a lower bound for our evidence precisely because the KL divergence term D[Q(z∣X)∣∣P(z∣X)]≥0, which implies logP(X)≥RHS.</li> <li>We cannot directly optimize logP(X), but we can do the next best thing which is to optimize the tractable RHS, given an appropriate choice of Q.</li> <li>As we increase the capacity of Q, the "error" term should become smaller and smaller, so the RHS will more accurately estimate the evidence (and lead to better optimization)</li> </ul> Notice how the RHS now resembles an auto-encoder: <ul> <li>Q "encodes" X into a latent z</li> <li>P "decodes" z back to reconstruct X</li> </ul> <h3 id="optimizing-the-objective"><a class="header" href="#optimizing-the-objective">Optimizing the Objective</a></h3> Now we need to perform SGD on the RHS. First we need to specify Q(z∣X). The usual choice is: Q(z∣X)=N(z∣μ(X;θ),Σ(X;θ)) Where μ and Σ are both neural networks with parameters θ that map a given X deterministically into a mean and variance vector respectively. We only need a vector for the variance because Σ is typically constrained to be a diagonal matrix. Because we chose both Q(z∣X),P(z) to be multi-variate gaussians, the KL divergence between them may now be computed in closed form (k is the dimensionality of the distribution): D[N(μ0,Σ0)∣∣N(μ1,Σ1)]=21(tr(Σ1−1Σ0)+(μ1−μ0)⊺Σ1−1(μ1−μ0)−k+log(detΣ0detΣ1))=21(tr(Σ(X))+μ(X)⊺μ(X)−k+logdet(Σ(X))) The second line above gives the general case, but since our P(z)=N(0,I), it reduces into the third line. The functions μ(X),Σ(X) express the fact that the parameters of the normal distribution of Q are determined by X. <h3 id="reparametrization-trick"><a class="header" href="#reparametrization-trick">Reparametrization Trick</a></h3> So we have the second term on the RHS expressed as a function of X, which can be optimized via SGD. What about the first term Ez∼Q[logP(X∣z)]? This term is more tricky, because it involves two steps. Suppose we approximate the expectation by performing SGD. Then we have to: <ul> <li>Sample a z∼Q</li> <li>Compute logP(X∣z) using the decoder</li> </ul> The first sampling step is not an operation that can be backpropagated through, so we cannot optimize this equation as-is. This is where the re-parametrization trick comes into play. The re-parametrization trick essentially moves the stochasticity of the sampling step out of the model forward pass into the data layer. Instead of sampling z∼Q directly, we first sample an intermediate ϵ∼N(0,I) and treat it as data. Then, we compute deterministically z=μ(X)+Σ1/2(X)×ϵ. This achieves the same effect as the direct sampling approach, but the difference is that the parameters of Q have entered the equation in a deterministic way that can be backpropagated. With the trick, we now have fully specified the optimization objective. The equation we take the gradient of is: EX∼D[Eϵ∼N(0,I)[logP(X∣z=μ(X)+Σ1/2(X)×ϵ)]−D[Q(z∣X)∣∣P(z)]] On the second term D[Q(z∣X)∣∣P(z)]: <ul> <li>Note that we have already spelt out the closed form expression above, which involves the encoder neural networks μ(X),Σ(X).</li> </ul> On the first term logP(X∣z): <ul> <li>This is the reconstruction loss which measures how well we reconstructed the input after passing it through the encoder back to the decoder again, so it involves both the encoder and decoder.</li> <li>Specifically, we first sample an ϵ∼N(0,I), then compute z=μ(X)+Σ1/2(X)×ϵ, which gives us one sample of an encoded input X.</li> <li>Recall above that we specified P(X∣z;θ)=N(X∣f(z;θ),σ2I), which is a PDF of a normal distribution with mean given by the decoder applied to sample z and some constant variance.</li> <li>Hence logP(X∣z)=−2σ21∣∣X−f(z;θ)∣∣22+const. which is an l2 distance</li> </ul> <h3 id="test-time-inference"><a class="header" href="#test-time-inference">Test Time Inference</a></h3> At test time, when we want to generate new samples, we simply sample a new z∼N(0,I), then feed it into the decoder to get a new X=f(z). Suppose we want to evaluate the probability of this new test example, i.e. P(X). This is not tractable for the reasons we discussed earlier, as it involves an integral over the distribution of z. However, we can make use of the ELBO concept and use the RHS of the ELBO equation as an approximation to P(X). There is still an expectation over z, but because sampling z∼Q gives an expectation that converges much faster than sampling z∼N(0,I), we can get a good sense of the probability by sampling a few times. <h2 id="extra-info"><a class="header" href="#extra-info">Extra Info</a></h2> This section tackles 3 questions to help our understanding: <ol> <li>How much error is introduced by the additional term D[Q(z∣X)∣∣P(z∣X)]?</li> <li>How is the VAE framework linked to Minimum Description Length?</li> <li>Do VAEs have regularization parameters analogous to sparsity penalties?</li> </ol> <h3 id="q1-error-from-lower-bound"><a class="header" href="#q1-error-from-lower-bound">Q1. Error from Lower Bound</a></h3> Given that we are optimizing for the RHS and not directly for logP(X), how much error does the additional term D[Q(z∣X)∣∣P(z∣X)] introduce? Since we assumed that Q(z∣X) takes the form of a high dimension gaussian, P(z∣X) must also take the form of a gaussian for the KL divergence term to go to 0. However, this is not necessarily the case - we make no assumption on the distribution of P(z∣X). The hope is that if f is sufficiently high-capacity, then there exists some f which both (i) maximizes P(X) and (ii) results in a gaussian-like P(z∣X). If such a function exists, then our objective would find it due to the D[Q(z∣X)∣∣P(z∣X)] term. <h3 id="q2-minimum-description-length-interpretation"><a class="header" href="#q2-minimum-description-length-interpretation">Q2. Minimum Description Length interpretation</a></h3> (I don't understand this part.) Another way to look at the RHS of the ELBO equation is in terms of information theory. logP(X) may be seen as the total number of bits required to construct a given X under our model using an ideal encoding. The RHS views this as a two-step process to construct X. <ul> <li>We first use some bits to construct z. D[Q(z∣X)∣∣P(z)] may be viewed as the expected information required to convert an uninformative sample from P(z) to a sample from Q(z∣X)</li> <li>In the second step, P(X∣z) measures the amount of information required to reconstruct X from z under an ideal encoding.</li> </ul> Hence the total number of bits −logP(X) is the sum of these two steps, minus a penalty we pay for Q being a sub-optimal encoding D[Q(z∣X)∣∣P(z∣X)]. <h3 id="q3-regularization-effect"><a class="header" href="#q3-regularization-effect">Q3. Regularization Effect</a></h3> It is interesting to view the D[Q(z∣X)∣∣P(z)] term as regularization, since it is encouraging our Q distribution to be similar to a simple distribution. In a usual sparse autoencoder, we have a λ parameter in an objective function: ∣∣ϕ(ψ(X))−X∣∣2+λ∣∣ψ(X)∣∣0 That is, for encoder ψ and decoder ϕ, we encourage the encoding to be sparse. Similarly, the KL divergence term encourages our encoder Q to be simple. Where does a similar parameter like λ enter into the ELBO equation? Recall that we chose a normal distribution for P(X∣z)∼N(f(z),σ2×I). It turns out that σ plays a similar role to λ, as we shall see. Using the PDF of the normal distribution, we have that logP(X∣z)=C−21∣∣X−f(z)∣∣2/σ2, where C is a constant that does not depend on f and can be ignored during optimization. In the ELBO equation, σ appears in the first term of the RHS but not the second term. Hence, by varying σ, we can control the relative weighting between the two terms. Specifically, a lower σ implies less regularization and a larger σ implies more regularization. <h2 id="conditional-vaes"><a class="header" href="#conditional-vaes">Conditional VAEs</a></h2> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="bateni-2017---affinity-clustering"><a class="header" href="#bateni-2017---affinity-clustering">Bateni 2017 - Affinity Clustering</a></h1> <a href="https://papers.nips.cc/paper_files/paper/2017/file/2e1b24a664f5e9c18f407b2f9c73e821-Paper.pdf">Bateni 2017 - Affinity Clustering: Hierarchical Clustering at Scale</a> This paper proposes a hierarchical clustering algorithm which they call "affinity clustering". It is essentially Boruvka's algorithm but with slight modifications, and the main contribution of the paper is a distributed algorithm to perform such clustering at scale. Much of this paper is devoted to theoretical analysis, but my interest here is just in implementing and understanding why the distributed algorithm is correct. <h2 id="naive-algorithm"><a class="header" href="#naive-algorithm">Naive Algorithm</a></h2> Suppose we have an undirected graph G=(V,E), where V is the set of nodes and E is the set of undirected edges. The naive algorithm for affinity clustering (closely following Boruvka's algorithm) is as follows: <ul> <li>Start with every node in its own cluster ci, where i=1,...,∣V∣.</li> <li>In each round, find the lowest weighted outgoing edge from each cluster ci and add the edge <ul> <li>i.e. for each cluster ci, find argmaxu,v {weight(u,v):u∈ci,v∈/ci}</li> <li>Note that this step ensures that the number of clusters at least halves, since each cluster is connected to another cluster</li> </ul> </li> <li>If the number of clusters becomes lower than k, we undo the most recent edge added until we get k clusters</li> <li>At the end of each round, we have the desired k number of clusters</li> <li>The next round then commences with these obtained clusters</li> </ul> Note that this is essentially Boruvka's algorithm (minus the undo steps), which is guaranteed to find the Minimum Spanning Tree (MST) of the graph G. This implies that if we start with G′:=MST(G), we can run the naive algorithm on the MST G′ and get exactly the same clusters as running it on G, because only the edges in G′ come into play when running affinity clustering on G. Since the number of nodes in the G′ is ∣V∣−1 (by definition of an MST), we would be able to do this efficiently on a single machine (unless |V| is much larger than a few millions). Therefore, we just need a distributed algorithm to find the MST of G efficiently, and we can perform affinity clustering efficiently. <h2 id="efficient-mst-algorithm"><a class="header" href="#efficient-mst-algorithm">Efficient MST algorithm</a></h2> The main algorithm is the distributed MST algorithm (see below). Some notations: <ul> <li>m is the number of edges ∣E∣, n is the number of nodes ∣V∣</li> <li>c measures the average density of the graph, i.e. the average number of edges per node. logn(m/n) is taken such that c∈[0,1], where c=0 implies that ∣V∣=∣E∣ and c=1 implies that G is a fully connected graph.</li> <li>0<ϵ<c is a density parameter controlling the final number of edges remaining before we run MST on a single machine. We can probably set it to 0.1 or something like that. A higher ϵ implies that we can run ReduceEdges for less steps, but we need more memory to run the final round of MST.</li> </ul> <div class="table-wrapper"><table><thead><tr><th style="text-align: center"><img src="papers/../images/bateni_2017_mst_algorithm.png" alt="Distributed MST Algorithm" /></th></tr></thead><tbody> <tr><td style="text-align: center">Distributed MST Algorithm</td></tr> </tbody></table> </div> The algorithm is really quite simple. The main idea is that we can independently process each subgraph of G comprising a random pair of vertex sets Vi⋃Uj in a distributed fashion. Each worker finds the MST of the subgraph Gi,j assigned to it, and any edge that is in E(Gi,j) but not in MST(Gi,j) may be removed from the global edge set E(G). In this way, we can whittle down the number of edges in E(G) until it is a small set that can fit in memory (i.e. in the order of O(n1+ϵ) which is not much larger than the number of nodes). Then we can just run MST one final time on the master node. So the algorithm really hinges on Lemma 6 in the paper, which tells us that removing edges in this distributed way will still give us the correct MST at the end. <div style="margin-left: 2em; border: 2px solid grey; padding: 10px;"> Lemma 6. Let G′=(V′,E′) be a (not necessarily connected) subgraph of the input graph G. If an edge e∈E′ is not in the MST of G′, then it is not in the MST of G either. My proof by contradiction. Suppose an edge e∈E(G) exists between nodes A,B. Let G′ denote a subgraph containing nodes A,B, and suppose for contradiction that e∈MST(G) but e∈/MST(G′). First cut MST(G) by removing the edge e such that A,B are in different partitions PA,PB (each partition is a set of nodes). Observe that since e exists in MST(G), it must be the lowest weight edge connecting PA and PB, since otherwise we could have replaced e with a lower weight edge to complete the MST. Now consider the subgraph G′ and partition the nodes in V′ according to PA,PB to form PA′⊂PA and PB′⊂PB. Consider the MST of G′ (might be a Minimum Spanning Forest instead, if not all nodes can be connected), and observe that there must exist a path between A and B in G′. Now remove any edges that cross PA′ and PB′ (call this edge set ER). Then add back all the edges in ER that do not connect a path between nodes A and B. Now there must exist exactly one edge remaining in ER that connects a path from A to B, since (i) a path exists between A and B in G and (ii) there cannot be more than one edge that does so, otherwise there would have been a cycle in MST(G′). We also know that this edge is not e. But e is the minimum weight edge between PA and PB, and therefore it must also be the minimum weight edge between PA′ and PB′. Hence this other edge could not have been in the MST of G′. We reach a contradiction. </div> Thus we are justified in removing edges in this way. This lemma is great because it allows us to independently process each subgraph, and if memory is of concern, we can also batch the number of subgraphs processed at each step according to the number of workers we have. Lastly, note that c is a dynamic parameter in the algorithm that measures the density (i.e. number of edges relative to number of nodes) of the graph at each step. Since we reduce the density of the graph in every step, c is stepped down progressively. This also results in k being stepped down progressively, where k controls the number of subgraphs at each step. We start with a large k, which requires less memory since small subgraphs implies that many edges are "chopped off". As the graph becomes less dense, we can afford to lower k progressively and remove more edges. <h2 id="implementation"><a class="header" href="#implementation">Implementation</a></h2> I can't seem to find an implementation of this algorithm, but it is probably easy to write a naive version of it using <code>multiprocessing</code> in python. We could store the edges in a <code>sqlite</code> database and distribute a batch of subgraphs at each turn, collect the edges to remove in the master and remove them from the db, and repeat. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="guo-2017---deepfm"><a class="header" href="#guo-2017---deepfm">Guo 2017 - DeepFM</a></h1> <a href="https://arxiv.org/pdf/1703.04247.pdf">Guo 2017</a> is an important innovation for recommender systems. The task tackled in this paper is to predict clickthrough rate (CTR) of recommended items. The paper argues that it is important to learn implicit interaction signals between user and item features. Suppose we define the <order> of a signal as the number of features involved in deriving the signal. Then: <ul> <li>Recency of an item is an <code>order-1</code> signal</li> <li>Item category and timestamp is an <code>order-2</code> signal, since e.g. demand for food items spikes at meal times</li> <li>Gender, game category and user age is an <code>order-3</code> signal, since e.g. male teenagers prefer shooting games</li> </ul> One can see that both low and high order signals are important for modelling CTR, and manual feature engineering is cumbersome to derive all such interaction rules. Hence we would like to have a deep learning model that can model such signals directly from raw user / item features and interaction data. This paper proposes to use Factorization Machines to model low order feature interactions and deep neural network to model high order feature interactions. <h2 id="setup-4"><a class="header" href="#setup-4">Setup</a></h2> Suppose we have a dataset comprising <code>n</code> instances of tuples (χ,y), where χ is an m-fields raw feature record where each field could be categorical or numerical. Each categorical field is represented as a vector of one-hot encoding, and each continuous field is represented as-is. Let xfieldj denote the vector representation of field <code>j</code> (where dimension of numericals is <code>1</code> and dimension of categoricals is the number of categories), and x denote the flattened vector of the xfieldj laid out horizontally. <h2 id="deepfm"><a class="header" href="#deepfm">DeepFM</a></h2> DeepFM comprises two distinct models: the <code>FM</code> component and the <code>deep</code> component which are simply summed together: y^=σ(yFM+ydeep). We go into each component below. The FM component is a factorization machine that captures order-1 and order-2 interactions between the raw features. We first project each feature field xfieldj to a <code>k</code> dimensional latent vector using a learned dimfieldj×k embedding matrix (in the paper, the authors set <code>k=10</code>). The latent vector representation of field <code>j</code> is denoted as Vj∈Rk. We compute the FM output as follows: yFM=⟨w,x⟩+i=1∑mj=i+1∑m⟨Vi,Vj⟩ The first term represents the order-1 representation of the features as-is. The second term is a pairwise inner product between the embedding representations of each feature field, which represents order-2 interactions between the features. The deep component tries to capture higher-order interactions between feature fields. This is done simply by laying out the embedding vectors Vj horizontally, such that we form a flat input vector of size <code>k \times m</code>, and we call this a(0). This fixed size vector is then fed into a multi-layer perceptron (or dense neural network) to finally return a sigmoid output. The standard forward equation for layer l of the MLP is denoted below. Note that the embedding layer is shared between the deep and FM networks, allowing the deep component to benefit from the faster learning of the FM component. a(l+1)=σ(W(l)a(l)+b(l))ydeep=σ(W(H+1)a(H)+b(H)) In the paper, the MLP is of the form <code>400</code>-><code>400</code>-><code>400</code>. Dropout is set to <code>0.5</code>. <h2 id="implementation-1"><a class="header" href="#implementation-1">Implementation</a></h2> Huawei has an <a href="https://github.com/reczoo/FuxiCTR/tree/main/model_zoo/DeepFM/DeepFM_torch">implementation of DeepFM</a> (amongst other models) in pytorch. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="van-den-oord-2017---vq-vae"><a class="header" href="#van-den-oord-2017---vq-vae">Van Den Oord 2017 - VQ-VAE</a></h1> <a href="https://arxiv.org/abs/1711.00937">Neural Discrete Representation Learning</a> This paper proposes a Variational Auto Encoder (VAE) that uses discrete (categorical) latent variables instead of continuous ones. Discrete latents are more desirable from an interpretability perspective, but historically have not been able to perform as well as continuous ones. This paper is apparently the first to bridge this performance gap. We are in the domain of computer vision, where the VAE is used to generate new images from existing ones. Note that although VQ-VAE starts with the VAE framework, it is not strictly variational because of the deterministic operations within. It is more precise to call it a (deterministic) autoencoder with vector quantization regularization. <h2 id="vaes"><a class="header" href="#vaes">VAEs</a></h2> A quick recap on regular VAEs: <ul> <li>We start with some input x which is an image</li> <li>The encoder encodes the input using q(z∣x) into latent space</li> <li>Where z in this paper is a discrete latent variable</li> <li>P(z) is the prior distribution (a uniform categorical in this paper)</li> <li>P(x∣z) is the decoder that decodes a given latent back into input space (i.e. generates an image)</li> </ul> <h2 id="vq-vae"><a class="header" href="#vq-vae">VQ-VAE</a></h2> In VQ-VAE, the prior and posterior distributions are categorical. Drawing from the categorical distribution gives us an index 1,...,K which is used to index into an embedding table comprising D dimensional embeddings. This extracted embedding is then used to represent the sample and fed into the decoder model. More specifically, let us define a latent embedding space e∈RK×D. That is, we have K discrete latents and D is the embedding dimension of each latent. Starting with an input x, we pass it into the encoder to produce ze(x), which is a grid of D-dimensional vectors. We then find the nearest embedding to ze(x) in the embedding table to get a categorical index k and a corresponding embedding ek for each encoded vector. <blockquote> Note that using our running example of image generation, ze(x) is encoded into a 2D grid of latents (say <code>32 x 32 x D</code>). We find the nearest embedding at each position such that we end up with a grid of <code>32 x 32</code> codebook indices. Since each position is discretized independently of the others, in the exposition we refer to x,ze(x) and so on as though it is one vector. </blockquote> One way to think about this operation is that we are applying a particular non-linear operation that maps ze(x)↦ek. Noticeably, this non linear operation is non-differentiable, which we will need to tackle later on. We can thus define the posterior probability distribution for q(z∣x) as: q(z=k∣x)={10if k=argminj∣∣ze(x)−ej∣∣2otherwise Note that this posterior distribution is deterministic. If we define a simple uniform prior for p(z), we get that the KL divergence is constant: DKL(q∣∣p)=logK. Recall that: DKL(Q∣∣P)=i∑Q(i)logP(i)Q(i) Since Q(i)=1 if i=k and 0 otherwise, only one term in the summation is non-zero: DKL(Q∣∣P)=1⋅log1/K1=logK Hence the KL divergence term is constant and ignored during training. <h2 id="learning"><a class="header" href="#learning">Learning</a></h2> At this point, we have fully specified the forward pass: <ul> <li>Start with input x</li> <li>Encode into ze(x)∈RD</li> <li>Use embedding table lookup to find nearest neighbour zq(x)=ek</li> <li>Decode into p(x′∣zq(x))</li> </ul> But it is not yet clear what the optimization objective and gradient flow should be. Recall that the standard VAE objective is to optimize: logP(X∣z)+DKL(Q(z∣X)∣∣P(z)) The first term is the reconstruction loss, where we draw from the Q distribution that maps an input x↦z (using the encoder plus some random noise). Then we decode the z back into input space and try to make it similar to input x. The second term is the KL divergence which tries to regularize the Q distribution to be as simple as possible. However, we cannot simply use this equation for the VQ-VAE: <ul> <li>While the reconstruction loss logP(X∣z) is no issue (we can use the standard gaussian formulation), the second KL divergence term is a constant as we saw above. Hence doing this naively just reduces to a standard deterministic auto-encoder.</li> <li>Another problem is that in computing P(X∣z), we have to go through the non-linear operation of looking up the embedding table. This does not allow the gradient to flow back into the encoder.</li> </ul> Hence the authors need to re-design the optimization objective. The first decision is to use the straight through estimator to circumvent the gradient issue. For the non-linear operation f(ze(x))=zq(x), we compute the forward pass normally (i.e. embedding table lookup) but simply pass the gradients through during backpropagation. This means that we approximate: dze(x)dL≈dzq(x)dL This allows the encoder to still receive gradient updates despite the non-differentiable operation. The theoretical justification for this operation is given in an earlier Bengio 2013 paper. Intuitively, if ze(x) is close to zq(x), the gradients should still be meaningful. The second decision is to use l2 distance to learn the embedding table. This is a form of dictionary learning. Specifically, we add a term to the loss: ∣∣sg[ze(x)]−e∣∣22 Note that: <ul> <li>e here refers to the closest embedding to a given ze(x). We want embeddings in the codebook to move toward the average encoded representation</li> <li>sg[] is the stop gradient operation (e.g. <code>.detach()</code> in pytorch). It uses the value of ze(x) but does not pass gradients back to the encoder. Since the objective of this loss term is to learn the codebook, we do not wish to pass gradients back to the encoder</li> </ul> The third decision is to add a commitment loss to bound the encoder outputs. This part feels a bit more arbitrary. The authors say that with just the first two terms, there is nothing that tethers the encoder output, which can grow arbitrarily and perpetually be far away from the codebook embeddings. The solution is to include the reverse direction from the dictionary learning loss: β∣∣ze(x)−sg[e]∣∣22 Notice that this is identical to the second term except that the stop gradient operator is applied to the codebook embedding. Thus this gradient pushes the encoder embedding to be closer to its nearest codebook embedding. β is a hyperparameter but the authors observed that results were robust to a wide range of β values (<code>0.1</code>to <code>2.0</code>). <blockquote> A natural question is to wonder why we need both the second and third term which are identical except for where the stop gradient is placed. Why can't we just do ∣∣ze(x)−e∣∣22 in a single term? It appears that this will result in unstable training, because both sides (encoder and codebook embeddings) are simultaneously moving. This is a common issue in training things like GANs. Separating the terms results in more stable training. </blockquote> <blockquote> There is also a close connection between VQ-VAE and k-means clustering (as a special case of expectation maximization). The step of assigning each encoder output to the nearest codebook embedding is akin to assign a cluster for each data point in k-means. The step of updating the codebook embedding is akin to updating the cluster centroids in k-means. This idea is explored in subsequent papers like <a href="https://arxiv.org/abs/1805.11063">Roy 2018</a>. </blockquote> Hence, the final training objective becomes: L=logp(x∣zq(x))+∣∣sg[ze(x)]−e∣∣22+β∣∣ze(x)−sg[e]∣∣22 <h2 id="evaluation-1"><a class="header" href="#evaluation-1">Evaluation</a></h2> The log likelihood of the complete model can be evaluated as the total probability over all latents k∈1,...,N. It is common practice to compute the log likelihood of a held out test set to evaluate how well our model has learned the data distribution: logp(x)=logk∑p(x∣zk)p(zk) Note that −log2p(x)/(H×W×Channels) is also computed to report bits per dimension, which is a common way to evaluate such VAE models on test data. This is literally the number of bits required to represent our data under this model. Because the decoder p(x∣z) is trained with z=zq(x) from MAP-inference, the decoder should end up placing all probability mass on p(x∣zq(x)) after full convergence and no probability mass on p(x∣z)∀z=zq(x). So we can write: logp(x)≈logp(x∣zq(x))p(zq(x)) <h2 id="learning-the-prior-distribution"><a class="header" href="#learning-the-prior-distribution">Learning the Prior Distribution</a></h2> Up to this point, we assumed that the prior P(z) is a uniform categorical for training the encoder, decoder and codebook embeddings. This may be viewed as a training trick or mathematical convenience to make our framework work. As you may recall, using the uniform prior resulted in a constant term for the KL-divergence, meaning that the term is ignored during training. At inference time, when generating a new image, using a uniform prior will result in incoherent images. Instead, we need to train a separate decoder like a pixelCNN or transformer over the latent space to generate a coherent latent grid for decoding. Specifically, we encode and quantize our training dataset into grids of <code>32 x 32</code> codebook indices. Then, an autoregressive decoder is trained to predict the codebook indices in a causal autoregressive way. For PixelCNN, it works on the 2D grid of <code>32 x 32</code>, but for transformer we need to linearize it in row-major form before training. Now at inference time, we can start with a uniform latent index for the first position, then generate the subsequent positions using the latent decoder. After we have generated a grid of <code>32 x 32</code> latents, we look up embeddings to get <code>32 x 32 x D</code> grid. This is passed to the decoder to get the generated image. <ul> <li>Note that since we are doing autoregressive generation in the compressed latent space, it is a lot faster than autoregressive generation in the original pixel or token space. This is the idea used in <a href="https://arxiv.org/abs/1803.03382">Kaiser 2018 - Fast Decoding in Sequence Models using Discrete Latent Variables</a>.</li> </ul> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="hamilton-2017---graphsage"><a class="header" href="#hamilton-2017---graphsage">Hamilton 2017 - GraphSAGE</a></h1> <a href="https://ar5iv.labs.arxiv.org/html/1706.02216">Paper</a>. This paper presents a framework to efficiently generate node embeddings, especially for previously unseen nodes. It calls this inductive learning, i.e. the model is able to generalize to new nodes, as opposed to previous frameworks which are transductive and only learn embeddings for seen nodes. For example, matrix factorization methods are transductive because we can only make predictions on a graph with fixed nodes, and need to be retrained when new nodes are added. GraphSAGE (i.e. Sample and Aggregation) aggregates feature information from the neighbourhood of a node to represent a given node. Feature information can include structural information (e.g. degree of a node) or other content-based information about a node, e.g. description of a job for a job node etc. <h2 id="setup-5"><a class="header" href="#setup-5">Setup</a></h2> We start with a graph G(V,E) that is provided by data. We have input features xv∀v∈V. For a user-chosen depth of K, we have K aggregator functions AGGk and K weight matrices Wk. We also have a neighbourhood function N:V→2V, which means it maps from a node v∈V to a set of nodes in V. Note that 2V is the powerset of the set V, i.e. the set of all possible subsets of V. We wish to generate low-dimensional, dense vector representations of each node zv. <h2 id="forward-propagation"><a class="header" href="#forward-propagation">Forward Propagation</a></h2> The algorithm for the forward propagation (in words) is as follows: <ol> <li>We start with hidden representations hv0←xv∀v∈V, i.e. at layer 0, we just use the input features to represent each node</li> <li>At depth k=1, we perform a neighbourhood aggregation step at each node v: hN(v)k←AGGk({huk−1:u∈N(v)})</li> <li>The aggregated vector is then passed through a dense layer to get the hidden representation at depth k. Note that σ represents a non-linear activation, such as <code>ReLU</code>: hvk=σ(Wk⋅CONCAT(hvk−1,hN(v)k))</li> <li>We L2-normalize each vector hvk ∀ v∈V</li> <li>We then repeat this process repeatedly for depths k=1,...,K</li> <li>We then take the last layer: zv←hvk</li> </ol> The intuition behind the forward propagation is that we use the neighbours of v to represent it. Importantly, we also include the hidden representation of the current node from the previous depth (analogous to a residual connection). At each depth level k, we increasingly pull more information from further reaches of the graph. Note that in the aggregation step, a subset of each node's neighbours are sampled uniformly (as opposed to taking the full neighbour set) to control the complexity of the algorithm. <h2 id="loss"><a class="header" href="#loss">Loss</a></h2> To train the weights, we define an unsupervised loss based on how well the embeddings are able to reconstruct the graph. Specifically, we have a loss which: <ul> <li>Rewards positive pairs for having a high dot product</li> <li>Penalizes negative pairs (vn being sampled negatives according to the negative sampling distribution Pn)</li> </ul> JG(zu)=−log(σ(zuTzv))−Q⋅Evn∼Pn(v)log(σ(−zuTzvn)) Alternatively, we can also define a supervised loss based on classification cross entropy loss, with presumably some form of negative sampling. The authors did not elaborate on this. <h2 id="aggregation-methods"><a class="header" href="#aggregation-methods">Aggregation Methods</a></h2> The authors explored a few ways to define the AGG function to aggregate neighbour embeddings together: <ul> <li><code>GraphSAGE-mean</code>: The element-wise mean of the neighbour embeddings is taken</li> <li><code>GraphSAGE-GCN</code>: Same as above, except that the current node's hidden representation from the previous depth hvk−1 is not included. The experiments show that omitting this residual connection actually leads to significant performance degradation.</li> <li><code>GraphSAGE-LSTM</code>: An LSTM is fitted over the sequence of embeddings. Since there is no inherent order to the neighbours, the authors randomize the ordering for each training sample</li> <li><code>GraphSAGE-pool</code>: An additional linear layer is added over the sequence of embeddings, before an element-wise <code>max-pool</code> operation is carried out</li> </ul> Generally from the experiments, it seems that <code>GraphSAGE-mean</code> is sufficient. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="ma-2018---entire-space-multi-task-model"><a class="header" href="#ma-2018---entire-space-multi-task-model">Ma 2018 - Entire Space Multi-Task Model</a></h1> <a href="https://arxiv.org/pdf/1804.07931.pdf">Ma 2018</a> tackles the problem of building a post-click Conversion Rate (CVR) prediction model. Note that CVR is the task of predicting conversions from impressions, whilst click-through rate prediction (CTR) is predicting clicks from impressions. In a typical recommender system, users follow the sequential process of <code>impression</code> -> <code>click</code> -> <code>conversion</code>, where conversion may refer to making a job application, purchase action etc. Usually, CVR models are built the same way as CTR models: a dataset of clicked impressions is prepared, and the converted items are labelled as relevant. A model is trained on this dataset and then used to make conversion predictions on all impressions. This paper argues that there are two problems with this approach: <ol> <li>Sample Selection Bias (SSB). The distribution of the training set (comprising only of clicked impressions) differs greatly from the distribution of the testing set (comprising all impressions), and this distribution shift will hurt generalization performance of the trained model.</li> <li>Data Sparsity (DS). The dataset for CVR (clicked impressions) is typically much less than the dataset for CTR (all impressions), and the lack of data makes model fitting difficult. The paper estimates that CVR dataset is typically <code>4%</code> of that of CTR dataset.</li> </ol> <h2 id="setup-6"><a class="header" href="#setup-6">Setup</a></h2> Denote the observed dataset to be S={(xi,yi→zi)}∣i=1N, with each sample tuple (x,y,z) representing one impression drawn from a distribution D with domain X×Y×Z. X is the feature space, and Y,Z are label spaces (i.e. <code>0</code> or <code>1</code>). Each feature vector x∈Rd captures all the user attributes, item attributes or user-item interaction for the impression event. The notation y→z represents the sequential nature where a click event y must precede a conversion event z. We can denote the various prediction tasks as follows: <ul> <li>Post-view clickthrough: pCTR=p(y=1∣x)</li> <li>Post-click conversion: pCVR=p(z=1∣y=1,x)</li> <li>Post-view click + conversion: pCTCVR=p(z=1,y=1∣x)=pCTR×pCVR</li> </ul> The conventional way of modelling <code>pCVR</code> is to construct a sample from only click impressions, i.e. Sc={(xj,zj)∣yj=1}, where clicked but not converted impressions are treated as negative samples. We can see than Sc⊂S. As mentioned above, there are problems with this approach. <h2 id="esmm-model"><a class="header" href="#esmm-model">ESMM Model</a></h2> The ESMM model breaks down the <code>pCTCVR</code> task into its constituents, and uses two neural networks to model <code>pCTR</code> and <code>pCVR</code> simultaneously. Based on the diagram, it seems to embed each user field and item field into a fixed-size embedding, where the user field embeddings are summed up element-wise to produce an overall user embedding. The same is done to produce an overall item embedding. The user and item embeddings are then concatenated together, and this combined embedding is fed into a dense layer to finally output a real score representing either <code>pCVR</code> or <code>pCTR</code>. The two scores are then multiplied together to form the final prediction of <code>pCTCVR</code>. Importantly, the projection (or lookup) layer from raw features to embedding is shared between the two neural networks. This allows the <code>pCVR</code> network in particular to benefit from the richer sample data that the <code>pCTR</code> network enjoys and addresses the data sparsity issue. <div class="table-wrapper"><table><thead><tr><th style="text-align: center"><img src="papers/../images/esmm_diagram.png" alt="ESMM Architecture" /></th></tr></thead><tbody> <tr><td style="text-align: center">ESMM Model Architecture (Figure 2 from ESMM Paper)</td></tr> </tbody></table> </div> Finally, the model is trained with a multi-task objective. Specifically, the losses are computed on the dataset with all impressions. The output <code>pCTR</code> is compared against clicks using a cross-entropy loss, and the output <code>pCTCVR</code> is compared against conversions using a cross-entropy loss. This multi-task loss allows us to exploit the sequential nature of the data generation process, such that fcvr only needs to model the delta aspect that leads from a click to a conversion. L(θcvr,θctr)=i=1∑Nl(yi,fctr(xi))+l(zi,fctr(xi)×fcvr(xi)) The authors show that modelling in this multi-task manner outperforms a setup where two models are trained independently to predict <code>CTR</code> and <code>CVR</code> respectively, and their product is taken to estimate <code>pCTCVR</code>. Unfortunately, we cannot replicate this joint-task learning setup with gradient tree-based models, at least not naively. <h2 id="details"><a class="header" href="#details">Details</a></h2> The authors set the dimension of each embedding vector to be <code>18</code>, and each MLP is <code>360</code> -> <code>200</code> -> <code>80</code> -> <code>2</code> dimensional. Adam solver is used with β1=0.9,β2=0.999,ϵ=10−8. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="kang-2018---sasrec"><a class="header" href="#kang-2018---sasrec">Kang 2018 - SASRec</a></h1> <a href="https://arxiv.org/abs/1808.09781">Self-Attentive Sequential Recommendation</a> This paper uses a transformer model with self causal attention to perform recommendations, by representing each user as a sequence of item embeddings and predicting the item interacted at time <code>t+1</code> by using all the information up to time <code>t</code>. <h2 id="background"><a class="header" href="#background">Background</a></h2> This paper came out shortly after the transformer (Attention is all you need) was invented. Up to this point, sequential recommendation was performed using Markov Chain methods or RNN-based methods. Since the self-attention mechanism of transformers is well suited to sequential modelling, this paper makes the natural adaptation of self-attention to the recommendation task. <h2 id="setup-7"><a class="header" href="#setup-7">Setup</a></h2> In the setting of sequential recommendation, we have for each user a sequence of item interactions Su=(S1u,S2u,...,S∣Su∣u) where each element represents an item. For computation reasons we may choose to truncate Su to the most recent n interactions. For simplicity we may also denote Su=(s1,...,sn). We have user and item sets U,I. Let us also define: <ul> <li>M∈R∣I∣×d as the full item embedding matrix with latent dimension d</li> <li>P∈Rn×d as the learned position embedding matrix</li> </ul> For each user, we receive Su and truncate it to the most recent n items. If there are less items, we left-pad the sequence with a constant zero vector. This results in an input embedding matrix E∈Rn×d for the user. Analogous to the language modelling task, the targets for each user is simply Su shifted to the left by one. In other words, the target at time step t would be the item interacted with at time step t+1. <h2 id="model-1"><a class="header" href="#model-1">Model</a></h2> Position Embeddings. We start by adding position embeddings to the user representation; absolute position embeddings are used here. Since this is a transformer model, the model has no notion of the item sequences if we do not inject the position embedding, and would not be able to learn that more recent items contain more valuable information about the next item to predict. The authors later show that visualizing the self-attention heatmap reveals that without position embedding, all items are attended to similarly, but with position embedding the attention weights are concentrated near the diagonal, i.e. more recent items are attended to stronger. Specifically, we simply add the position embedding matrix to the input embedding matrix, such that: E^t=Mst+Pt Self attention. The standard scaled dot product attention is used to perform self attention on the input embedding. Specifically: SA(E^)Attention(Q,K,V)=Attention(E^WQ,E^WK,E^WV), where=softmax(d<svg xmlns="http://www.w3.org/2000/svg" width='400em' height='1.08em' viewBox='0 0 400000 1080' preserveAspectRatio='xMinYMin slice'><path d='M95,702 c-2.7,0,-7.17,-2.7,-13.5,-8c-5.8,-5.3,-9.5,-10,-9.5,-14 c0,-2,0.3,-3.3,1,-4c1.3,-2.7,23.83,-20.7,67.5,-54 c44.2,-33.3,65.8,-50.3,66.5,-51c1.3,-1.3,3,-2,5,-2c4.7,0,8.7,3.3,12,10 s173,378,173,378c0.7,0,35.3,-71,104,-213c68.7,-142,137.5,-285,206.5,-429 c69,-144,104.5,-217.7,106.5,-221 l0 -0 c5.3,-9.3,12,-14,20,-14 H400000v40H845.2724 s-225.272,467,-225.272,467s-235,486,-235,486c-2.7,4.7,-9,7,-19,7 c-6,0,-10,-1,-12,-3s-194,-422,-194,-422s-65,47,-65,47z M834 80h400000v40h-400000z'/></svg>QKT)V Where WQ,WK,WV∈Rd×d are the projection matrices. We then make sure to mask the softmax attention matrix in a causal manner, so that there can be no interaction between Qi and Kj for all j>i. Feedforward network. A point-wise two-layer feedforward network is applied to the output of the self attention (i.e. S∈Rn×d), like so: FFN(S)=ReLU(SW(1)+b(1))W(2)+b(2) Where W(1),W(2)∈Rd×d and b(1),b(2)∈Rd. Note that in the feedforward networks, there remains no interaction between any Si,Sj at different time positions. Stacking self attention layers. Now we stack the self attention layers and also apply (i) residual connection, (ii) dropout and (iii) LayerNorm each time. This is standard practice in transformers and leads to more stable training. S(b)F(b)=g(SA)(F(b−1))=g(FFN)(S(b)) Where we define the composite function g(SA) as follows, and g(FFN) is defined similarly. g(SA)(X):=X+Dropout(SA(LayerNorm(X))) <blockquote> Note: In modern transformers, the LayerNorm is replaced by the simpler RMSNorm and the ReLU function is replaced by the geLU function. </blockquote> This gives us the full specification for one layer (layer b) of the transformer. Several layers are stacked to provide the full model. Prediction. After the final B layer, we have F(B)∈Rn×d as our representation. The predicted score at each time step t for any item i∈I is made according to a simple dot product: ri,t=Ft(B)MiT <h2 id="training-loss"><a class="header" href="#training-loss">Training Loss</a></h2> As discussed above, the target for each time step t is simply the next item at time step t+1. Specifically, if we define ot as the target output at time step t, we have: <ul> <li>ot=<pad> if st is a padding item</li> <li>ot=st+1 for 1≤t<n</li> </ul> <blockquote> Note: Not sure if we need to predict the padding item, or just simply mask the loss at those positions. Similarly for time step n, where we do not know the next item to predict after the last item in the sequence. </blockquote> Finally, the binary cross entropy loss is chosen as the objective function for each time step t. Specifically, a random negative item that user u has not interacted with is sampled for each time step and used as the negative example. The loss is: L=−Su∑t∑<svg xmlns="http://www.w3.org/2000/svg" width='0.667em' height='3.600em' viewBox='0 0 667 3600'><path d='M403 1759 V84 H666 V0 H319 V1759 v0 v1759 h347 v-84 H403z M403 1759 V0 H319 V1759 v0 v1759 h84z'/></svg>log(σ(rot,t))+j∈/Su∑log(1−σ(rj,t))<svg xmlns="http://www.w3.org/2000/svg" width='0.667em' height='3.600em' viewBox='0 0 667 3600'><path d='M347 1759 V0 H0 V84 H263 V1759 v0 v1759 H0 v84 H347z M347 1759 V0 H263 V1759 v0 v1759 h84z'/></svg> <blockquote> Note: The binary cross entropy loss is chosen here with one sampled negative per time step. A later paper <a href="https://arxiv.org/abs/2309.07602">Turning Dross Into Gold Loss: is BERT4Rec really better than SASRec?</a> will show that using softmax cross entropy loss with many sampled negatives will lead to much better performance. </blockquote> <h2 id="experiments-1"><a class="header" href="#experiments-1">Experiments</a></h2> For the experiments, two attention layers were used (i.e. B=2). Item embeddings are shared between the input embedding layer and also used in the prediction layer (for Mi). The latent dimension is set to d=50. The ablation studies found that: <ul> <li>Increasing number of layers saturates at around B=3</li> <li>Using multi-head attention did not improve over single head attention</li> <li>The absolute position embeddings generally improved performance relative to no position embeddings</li> </ul> <h2 id="implementation-2"><a class="header" href="#implementation-2">Implementation</a></h2> The authors endorsed a pytorch implementation <a href="https://github.com/pmixer/SASRec.pytorch/blob/main/python/model.py">here</a>. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="kaiser-2018---fast-decoding-in-latent-space"><a class="header" href="#kaiser-2018---fast-decoding-in-latent-space">Kaiser 2018 - Fast Decoding in Latent Space</a></h1> <a href="https://arxiv.org/abs/1803.03382">Fast Decoding in Sequence Models using Discrete Latent Variables</a> This paper uses the idea in <a href="papers/./van_den_oord_2017.html">VQ-VAE</a> to do fast decoding in latent space. <h2 id="main-idea-1"><a class="header" href="#main-idea-1">Main Idea</a></h2> Transformers parallelize during training well, but are slow at inference due to autoregressive nature. This paper proposes to autoregressively generate in discrete latent space instead, which can be much faster depending on the compression: <ul> <li>First auto-encode target sequence into discrete latent space which is a shorter sequence</li> <li>Train a decoder to generate in discrete latent space autoregressively</li> <li>Decode output sequence from the shorter latent sequence in parallel</li> </ul> Here is an infographic from nano banana that captures the idea of generation in latent space: <div class="table-wrapper"><table><thead><tr><th style="text-align: center"></th></tr></thead><tbody> <tr><td style="text-align: center"><img src="papers/../images/kaiser_2018_generation_in_latent_space.jpg" alt="Generation in Latent Space" /></td></tr> <tr><td style="text-align: center">Fast Generation in Latent Space</td></tr> </tbody></table> </div> <h2 id="setup-8"><a class="header" href="#setup-8">Setup</a></h2> When generating sequential output, autoregressive models generates y1,...,yn in a canonical order. This is because the model is trained to predict: P(yt∣yt−1,yt−2,...,y1) During training, because ground truth is known, we can train in parallel. During decoding, this is a fundamental limitation. The proposal is to first encode the original sequence into a shorter sequence of discrete tokens l1,...,lm, where m is much smaller than n. This latent sequence is also autoregressively predicted, and subsequently decoded in parallel back into the original token space. In the experiments, m=8n. <h2 id="latent-transformer"><a class="header" href="#latent-transformer">Latent Transformer</a></h2> The latent transformer is described in terms of a machine translation task. Given an input sequence x in English, and a corresponding output sequence y in German, our input-output pair is (x,y)=(x1,...,xk,...,y1,...,yn). The high level architecture is: <ul> <li>The function ae(y,x) will autoencode y into a shorter sequence l=l1,...,lm of discrete latent tokens using some discretization technique</li> <li>The latent prediction model lp(x) which is a transformer will autoregressively predict l based on x</li> <li>The decoder ad(l,x) is a parallel model that will decode y from l and the input sequence x</li> </ul> <h3 id="loss-1"><a class="header" href="#loss-1">Loss</a></h3> The reconstruction loss compares the encoded - decoded output to the original target sequence y: y^Lrecon=ad(ae(y,x),x)=CrossEntropyLoss(y^,y) The latent prediction loss compares the true latents generated by the autoencoder to the generated latents lp(x): ltruelpredictedLlatent=ae(y,x)=lp(x)=CrossEntropyLoss(ltrue,lpredicted) The final loss is the sum of the two: L=Lrecon+Llatent <h3 id="autoencoder"><a class="header" href="#autoencoder">Autoencoder</a></h3> Now we describe each part in more detail. The diagram below summarizes the architecture of the autoencoder ae(y,x). <div class="table-wrapper"><table><thead><tr><th style="text-align: center"></th></tr></thead><tbody> <tr><td style="text-align: center"><img src="papers/../images/kaiser_2018_autoencoder_architecture.jpg" alt="Autoencoder Architecture" /></td></tr> <tr><td style="text-align: center">Autoencoder architecture</td></tr> </tbody></table> </div> The input to the autoencoder is of shape <code>length x hidden_size</code>. The aim is to shorten the sequence length by downsampling using convolutions. We first pass through a residual block which aims to learn local features using convolutions: <ul> <li>Pass through <code>relu</code></li> <li>Pass through 1D convolution with <code>k=3, s=1</code>. This means we look over a window of <code>3</code> tokens each time, and take a step of <code>1</code> token each step. <ul> <li>Since stride is <code>1</code>, the sequence length is unchanged</li> </ul> </li> <li>Pass through layer norm + add the residual connection</li> </ul> Then we pass through a standard self attention block. Then we pass <code>c</code> times through a downsampling convolution: <ul> <li><code>k=2</code> means that we look over a window of <code>2</code> tokens</li> <li><code>s=2</code> means that we stride over <code>2</code> steps at a time: this effectively halves the sequene length</li> <li>Doing this <code>c</code> times means we get C=2c reduction in length</li> </ul> Finally we pass through the <code>bottleneck</code> function, which has yet to be described. But essentially the <code>bottleneck</code> function converts the representation of <code>hidden_size</code> at each position into a discrete token, and then represents it instead by the embedding of the discrete token using a lookup. <h3 id="decoder"><a class="header" href="#decoder">Decoder</a></h3> The decoder ad(l,x) decodes from latent space back into token space. <div class="table-wrapper"><table><thead><tr><th style="text-align: center"></th></tr></thead><tbody> <tr><td style="text-align: center"><img src="papers/../images/kaiser_2018_decoder_architecture.jpg" alt="Decoder Architecture" /></td></tr> <tr><td style="text-align: center">Decoder architecture</td></tr> </tbody></table> </div> We first pass through <code>c</code> up-steps, each of which will double the sequence length of the encoded sequence: <ul> <li>The same residual block from before is used</li> <li>Self attention layer is used</li> <li>An up-conv step is performed: <ul> <li>A standard feed forward MLP is used that doubles the hidden dimension</li> <li>A reshape operation is done to translate the extra hidden dimension into sequence length instead</li> </ul> </li> </ul> Finally, the decompressed sequence of <code>length x hidden_size</code> is fed into a standard transformer decoder to decode the target sequence y. <ul> <li>Note that for the first 10k steps of training, instead of feeding the decompressed sequence, the authors pretrain this decoder transformer model by feeding the target sequence y (i.e. training it like a standard transformer decoder). This is to warm up the head so that gradients to upstream components are reasonable.</li> </ul> <h2 id="discretization-bottleneck"><a class="header" href="#discretization-bottleneck">Discretization Bottleneck</a></h2> This is actually the main part of the paper - how do we discretize our token sequence into latent space so that we can predict the latent tokens in an autoregressive manner? The <code>bottleneck</code> function takes a target token embedding y and maps it to a discrete token in [K]:={1,2,...,K}: <ul> <li>First, take y and encode it: enc(y)∈RD, where D is the dimension of the latent space.</li> <li>Next, pass the encoding through a bottleneck to product a discrete latent token zd(y)∈[K]</li> <li>Look up embedding table using zd(y) to produce an embedding zq(y)∈RD to pass to the decoder</li> </ul> Note that the <code>bottleneck</code> function is applied to each sequence position in parallel, indepedently. <h3 id="gumbel-softmax"><a class="header" href="#gumbel-softmax">Gumbel Softmax</a></h3> The gumbel-softmax trick is found in <a href="https://arxiv.org/abs/1611.01144">Jang 2016</a>. It is a popular discretization technique. First, we project the encoder output enc(y)∈RD using a learned projection matrix W∈RK×D, such that we get logits l: l=Wenc(y)∈RK We can simply extract the discrete code zd(y) for the decoder as: zd(y)=i∈[K]argmaxli For evaluation and inference, we can simply use j:=zd(y) to look up an embedding table e∈RK×D to get zq(y)=ej. However, we cannot use this approach for training as the argmax and embedding lookup are non-differentiable. For training, the Gumbel-softmax trick is used to make the whole thing differentiable. <ul> <li>First, we draw g1,...,gK∈R i.i.d samples from the standard Gumbel distribution gi∼−log(−logu), where u∼U(0,1)</li> <li>Then we compute the weight vector w∈RK using a softmax, where for each latent dimension i∈[K]: wi=∑i=1Kexp((li+gi)/τ)exp((li+gi)/τ)</li> <li>Now we simply compute the input to the decoder as the weighted sum of the embedding table: zq(y)=we∈RD</li> </ul> For low temperatures of τ, w would be close to representing the 1-hot arg-max index of l, which is what we use for evaluation and inference. Setting τ too high would cause divergence between training and inference, which is not ideal. <blockquote> Why do we add the gumbel gi? Ideally during training, we want the encoder logits Softmax(l) to represent probabilities of a categorical distribution, from which we sample the discrete code zd(y). This ensures that we have: <ul> <li>Exploration from the stochasticity</li> <li>Discretization. We also force the embedding zq(y) to be a single embedding (i.e. discretization) rather than a smooshed average of several embeddings.</li> </ul> The gumbel-max trick tells us that choosing k=argmax(li+gi) results in an index k distributed exactly according to P(class k) as dictated by the softmax of the logits, which is what we want. Since the gumbel-max is not differentiable, we simply relax it into a softmax operation, giving us the gumbel-softmax. </blockquote> <h3 id="vector-quantization"><a class="header" href="#vector-quantization">Vector Quantization</a></h3> The VQ-VAE is covered in <a href="papers/./van_den_oord_2017.html">VQ-VAE</a>. Here is a brief recap. <ul> <li>The encoder output enc(y)∈RD is used to do nearest neighbour lookup on embedding vectors e∈RK×D to give us: <ul> <li>zq(y)=ek and zd(y)=k</li> <li>Where k=argminj∈[K]∣∣enc(y)−ej∣∣2</li> </ul> </li> <li>The loss is the reconstruction loss + the commitment loss between encoder output and codebook vectors: L=lrecon+β∣∣enc(y)−sg(zq(y))∣∣2</li> </ul> It turns out that the gradient update can be replaced by an exponential moving average over: <ul> <li>The codebook embeddings ej; and</li> <li>The count cj measuring the number of times ej is selected as a nearest neighbour</li> </ul> For a given mini-batch of target embeddings {y1,...,yl,...}, we can update the exponential moving average for counts like so: cj←λcj+(1−λ)l∑1[zq(yl)=ej] And the codebook embeddings ej updated as the average encoder output selected for ej: ej←λej+(1−λ)l∑cj1[zq(yl)=ej]enc(y) λ is a decay parameter and set to 0.999 in the experiments Note. The paper also covers two other discretization techniques, namely improved semantic hashing and decomposed vector quantization (to counteract codebook collapse), but I don't cover it for now. <h2 id="results"><a class="header" href="#results">Results</a></h2> The results show that the discretization bottleneck is not as good as a standard transformer, but beats all previous discrete methods and runs faster than a normal transformer. Ablation results show that: <ul> <li>Increasing the codebook size K improves performance</li> <li>Increasing the ratio n/m (i.e. shortening the sequence more aggressively) leads to speedup in inference but also clearly hurts performance</li> </ul> Hence it seems like there is no free lunch here. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="reimers-2019---sentence-bert"><a class="header" href="#reimers-2019---sentence-bert">Reimers 2019 - Sentence-BERT</a></h1> <a href="https://ar5iv.labs.arxiv.org/html/1908.10084">Link to paper</a>. Sentence-BERT (or SBERT) was one of the first papers to suggest a way to fine-tune BERT models to generate useful embeddings that can be used for search / retrieval. Prior to SBERT, BERT models were mainly used for sentence pair regression tasks by passing two sentences into the transformer network and adding a classification head on top to produce a float value. We can call this the cross-encoder approach. In other words, researchers only cared about the final prediction and did not make use of the embeddings, or the final vector representation of the inputs. This approach is suitable for reranking a small number of documents but not for nearest neighbour search in a corpus with millions of documents. Naively, one can element-wise average the BERT embeddings at the final layer to produce a vector representation of the text. This vector representation can then be used for nearest neighbour search or clustering. However, because BERT was not explicitly trained for this objective, it results in rather bad sentence embeddings, often worse than GloVe embeddings. <h2 id="method"><a class="header" href="#method">Method</a></h2> The SBERT paper presents 3 different training objectives, all of which perform well on embedding similarity tasks. The choice of objective depends on the dataset: <ol> <li> Classification objective. This is for tasks where the objective is to predict a label given two sentences A, B. We pass each sentence into the BERT network and a pooling layer to get two vector representations, u and v. The pooling layer can be (i) take the <code>[CLS]</code> token embedding, (ii) take the element-wise mean or (iii) take the element-wise max. We then create a concatenated vector (u, v, ∣u−v∣) which is fed into a softmax classifer. The network is trained using cross-entropy loss. Note that this siamese approach (where each sentence is passed into the same network) differs a little from the typical cross-encoder approach, where the sentences are concatenated as a string with the <code>[SEP]</code> token before passed into the network. The latter approach is presumably more powerful because the attention mechanism can attend to all pairwise relationships </li> <li> Regression objective. This is for tasks where the objective is to predict a float given two sentences A, B. Given the vectors u and v, the cosine similarity is simply taken to generate a float between −1 and 1. The cosine similarity is then compared with the actual float value using mean-squared error to generate a loss. </li> <li> Triplet objective. This is for tasks where each data point is a triplet (anchor sentence a, positive sentence p, negative sentence n). We then minimize the loss function, where m is the margin: max(∣∣sa−sp∣∣−∣∣sa−sn∣∣+m,0) </li> </ol> <h2 id="ablation"><a class="header" href="#ablation">Ablation</a></h2> <ol> <li>Pooling strategy. Using <code>[CLS]</code> or mean seems to be largely similar. The authors saw some degradation using max for the regression objective.</li> <li>Concatenation. For the classification objective, the concatenation strategy makes some difference. In particular, using (u,v) yields ρ=0.66 but (u,v,∣u−v∣) yields ρ=0.81. Thus the element-wise difference is important in yielding useful embeddings, probably because it can be used to push similar sentences together and dissimilar sentences apart. The authors also found that adding element-wise multiplication u∗v does not help.</li> </ol> <h2 id="takeaway"><a class="header" href="#takeaway">Takeaway</a></h2> It is interesting that the classification objective, which is close to a cross-encoder framework, is also able to learn useful embeddings by adding the difference operation ∣u−v∣. This suggests that we can train a cross encoder and simultaneously get useful embeddings for nearest neighbour retrieval at the same time. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="yi-2019---logq-correction-for-in-batch-sampling"><a class="header" href="#yi-2019---logq-correction-for-in-batch-sampling">Yi 2019 - LogQ Correction for In Batch Sampling</a></h1> <a href="https://research.google/pubs/sampling-bias-corrected-neural-modeling-for-large-corpus-item-recommendations/">Yi 2019 - Sampling Bias Corrected Neural Modelling for Large Corpus Item Recommendations</a> This paper proposes a way to perform logQ correction for sampling bias introduced by in-batch negative sampling when training two tower models. The algorithm proposed is a streaming algorithm that estimates item frequencies based updates after seeing each mini batch. <h2 id="setup-9"><a class="header" href="#setup-9">Setup</a></h2> Let xi, yj denote a user and item respectively, where there are i=1,...,N users and j=1,...,M items. Let u(.) and v(.) denote user and item embedding functions that map each xi and yj to Rk. These functions are typically: <ul> <li>Some sentence transformer model for texts</li> <li>Some hash embedding in the collaborative filtering setting</li> </ul> The output of the model is the inner product of the embeddings, i.e. s(x,y)=<u(x),v(y)>. The goal is to train the model from a training dataset of T user-item interactions, denoted by T:={xi,yi,ri}i=1T, where xi, yi are the interacting query and item and ri∈R is the associated reward. <ul> <li>Typically ri=1 to denote an interaction</li> <li>We can also use ri to denote some quality weight, e.g. time spent on product</li> </ul> Given a query x, we typically model the conditional probability of picking item y based on the softmax function. θ parametrizes the embedding model: P(y∣x;θ)=∑j∈[M]es(x,yj)es(x,y) We then design the loss function as a weighted log likelihood of the training interactions: LT(θ):=−T1i∈[T]∑ri⋅log(P(yi∣xi;θ)) <h2 id="in-batch-sampling"><a class="header" href="#in-batch-sampling">In Batch Sampling</a></h2> In practice, the denominator for P above is not feasible to compute when the number of items M is very large. The common practice is to sample only a subset of B items that are drawn in a mini batch. Hence given a mini batch of B pairs {(xi,yi,ri)}i=1B and for any i∈[B], the batch softmax becomes: PB(yi∣xi;θ)=∑j∈[B]es(xi,yj)es(xi,yi) Note that each xi,yi refers to a positive pair. However, the batch softmax above is usually a very biased estimate of the full softmax. This is because our training data usually has a heavy bias toward popular items, hence the likelihood of a popular item being included in the denominator is usually quite skewed. In other words, our model trained with this biased likelihood function may have a low training loss against popular items in the denominator during training. But when used in retrieval, the model may be assigning high scores to rare items that should be negatives, just that our model did not have a chance to discriminate against them due to the biased sampling during training. This issues underlies the common phenomenon when training such retrieval embedding models where the reranking performance is good but retrieval performance is very bad. The reason is that reranking is often performed against popular items that the model sees often, but retrieval by definition searches across the whole item catalogue. Hence retrieval is (from this perspective) a harder task than reranking. Special attention must be paid during training to ensure that the model learns to discriminate well against all items in the catalogue, and this logQ correction is one of the methods at our disposal. In <a href="https://infoscience.epfl.ch/server/api/core/bitstreams/bf28b394-cbd0-4f4c-ba3b-0bc7577ad95e/content">Adaptive Importance Sampling to Accelerate Training of A Neural Probabilistic Language Model</a>, the authors propose the following way to correct the biased batch softmax by correcting each score logit: sc(xi,yj)=s(xi,yj)−log(pj) Where pj denotes the probability of sampling an item j in a random batch. With this correction, we can denote the batch softmax as: PBc(yi∣xi;θ)=esc(xi,yi)+∑j∈[B],j=iesc(xi,yj)esc(xi,yi) And finally we have the batch loss function as: LB(θ):=−B1i∈[B]∑ri⋅log(PBc(yi∣xi;θ)) <h2 id="estimating-sampling-probability-in-stream-setting"><a class="header" href="#estimating-sampling-probability-in-stream-setting">Estimating Sampling Probability in Stream Setting</a></h2> Notably, the batch loss function does not require holding a fixed set of items in memory to serve as negative candidates, making it suitable for use in a streaming training data setting. Thus, the authors propose a method to estimate the sampling probability pj in a streaming fashion as well. The first observation is that it is easier to track the number of steps (or batches) between two consecutive hits of item j. e.g. if we only get one item once every 50 batches, then p=0.02. The proposed algorithm is as follows: <ol> <li>Initialize Arrays A,B with size H</li> <li>Let h(.) be a hash function from an item ID to [H]</li> <li>At batch t, sample a batch of items. For each item y in the batch: <ul> <li>B[h(y)]←(1−α)⋅B[h(y)]+α⋅(t−A[h(y)])</li> <li>A[h(y)]←t</li> </ul> </li> <li>At inference time, the sampling probability for item y will be py=1/B[h(y)].</li> </ol> <h2 id="other-notes"><a class="header" href="#other-notes">Other Notes</a></h2> The authors note that adding l2-normalization to embeddings improves model trainability and leads to better retrieval quality. Also, adding a temperature τ to each logit helps to sharpen the predictions. In their experiment, the best τ is usually around <code>0.05</code> (i.e. logits get multipled by <code>20x</code>). <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="zhao-2019---recommending-what-to-watch-next"><a class="header" href="#zhao-2019---recommending-what-to-watch-next">Zhao 2019 - Recommending What to Watch Next</a></h1> <a href="https://daiwk.github.io/assets/youtube-multitask.pdf">Recommending What Video to Watch Next: A Multitask Ranking System</a> This paper covers Youtube's recommendation system for what video to watch next, given the current video. The main contributions of this paper are: <ul> <li>Handling position bias</li> <li>Multi-gated Mixture of Experts to handle multi-task objectives</li> </ul> <h2 id="background-1"><a class="header" href="#background-1">Background</a></h2> This paper focuses on the ranking task of next video recommendation. There are specific challenges to their setting: <ol> <li>Multiple signals. For each video recommendation, there are engagement signals such as clicked, duration watched etc., and also user satisfaction signals such as liked, shared etc. Optimizing for these tasks is potentially conflicting, but multi-task learning also has potential to learned shared representations across tasks.</li> <li>Position bias. Position bias can degrade recommender performance since it may learn signals irrelevant to relevance. Tackling this issue can lead to online performance gains.</li> <li>Large Scale. The industry scale of recommending from billions of videos means that training quality might need to be traded off for efficiency. For this reason, the paper opts for a neural network based point-wise ranking model (as opposed to pair-wise or list-wise training methods like InfoNCE). This is also why the paper focuses on an architecture that can efficiently share parameters between different input feature modalities and prediction tasks.</li> </ol> The candidate generation is performed using multiple retrieval models, each of which captures a different aspect. For example, one algorithm matches topics of the query video, whilst another is based on co-watch statistics. A sequence model similar to <a href="https://static.googleusercontent.com/media/research.google.com/en//pubs/archive/45530.pdf">Covington 2016 - Deep Neural Networks for YouTube Recommendations</a> is also used. The efficient gramian methods in <a href="https://arxiv.org/abs/1807.07187">Krichene 2018 - Efficient Training on Very Large Corpora via Gramian Estimation</a> was also used. For ranking, a few hundred candidates are retrieved and passed to the ranker. Point-wise ranking loss is chosen for simplicity and scaleability. <h2 id="mixture-of-experts"><a class="header" href="#mixture-of-experts">Mixture of Experts</a></h2> It is common for ranking systems to have multiple objectives, due to having multiple implicit feedback signals in a recommender system. There are some approaches to deal with this: <ul> <li>Explicitly modelling the relationship between the signals. E.g. in <a href="papers/./ma_2018.html">ESMM</a>, the conversion probability head is modelled as a multiplication of the click probability with a separate head. In this way, the conversion head only needs to learn the "delta" signals that cause a click to turn into a conversion. However, such an approach only works when the relationship between tasks is known, and is also not very scaleable when the number of tasks increases.</li> <li>Shared Bottom Layer. The typical architecture is to have a shared representation layer at the bottom (from the inputs), and then to have a task specific head for each task connected to this shared bottom (see figure below). However, such hard parameter sharing has been shown to sometimes harm the learning when correlation between tasks is low.</li> </ul> <div class="table-wrapper"><table><thead><tr><th style="text-align: center"><img src="papers/../images/zhao_2019_shared_bottom_architecture.png" alt="Shared Bottom Architecture" /></th></tr></thead><tbody> <tr><td style="text-align: center">Shared Bottom Architecture</td></tr> </tbody></table> </div> The mixture of experts architecture proposed is covered in <a href="https://dl.acm.org/doi/pdf/10.1145/3219819.3220007">Ma 2018 - Modeling Task Relationships in Multi-task Learning with Multi-gate Mixture of Experts</a>. The main idea is that we have several expert layers that branch off from the shared bottom, and we also have a gating head (ending in a softmax) for each expert tower. The representation for each task is thus a unique weighted combination (the weights coming from the gating heads) of the expert layers. This somehow allows the model to "distribute work" to different experts to modularize information for each task. Note that the number of experts can differ from the number of tasks. Specifically, focusing on task k, let us denote: <ul> <li>x∈Rd as the representation arising from the shared bottom layer</li> <li>fi(x)∈Rd as the ith expert where i=1,...,n. <ul> <li>In this paper, f is simply a sequence of ReLU layers.</li> <li>Here I denote the expert representation as having the same dimension d for simplicity but that does not need to be the case.</li> </ul> </li> <li>g(i)k(x)∈R is the gating weight for task k from the ith expert</li> <li>gk(x) is the vector form comprising n elements for n experts <ul> <li>It is obtained by gk(x)=softmax(Wk(x))</li> <li>Where Wk∈Rn×d is the gating weights for task k</li> </ul> </li> <li>hk is a final sequence of ReLU layers for task k</li> </ul> Starting from the shared bottom representation x, for each task k we perform: <ul> <li>Obtain the gating weight vector gk(x)=softmax(Wk(x))</li> <li>Obtain the weighted representation for task k as fk(x)=∑i=1ng(i)k(x)⋅fi(x)</li> <li>Obtain the final prediction logit for task k as yk=hk(fk(x))∈R</li> </ul> We see that this architecture is quite flexible in supporting an arbitrary number of experts and tasks. In contrast to <a href="papers/./ma_2018.html">ESMM</a>, we do not need any knowledge of how the specific tasks relate to each other to specify this network. The authors mention that in distributed training, there is a danger that the softmax gating network may experience an imbalanced expert distribution problem, where gating networks converge to have zero-utilization on most experts. They experienced this 20% of the time. To mitigate this issue, dropout is performed on the gating networks by setting the softmax weight for each gate to be 0 with probability <code>10%</code>. The softmax vector gk is then renormalized to sum to <code>1</code> after the dropout. With this approach, they completely eliminated the polarization issue. It is not clear whether this issue happens in single-machine training, but nonetheless the fix seems easy to implement and seems like a low-hanging fruit to do. <h2 id="position-bias"><a class="header" href="#position-bias">Position Bias</a></h2> A common problem with implicit feedback data is position bias, i.e. the propensity of users to click items in higher positions due to the position rather than the relevance of the item. There are some ways to deal with this: <ul> <li>Inject the position as an input feature in model training to learn the position bias, and then set this feature to a fixed constant when serving</li> <li>Learn the position bias explicitly from offline data or random experiments, then use methods like Inverse Propensity Scoring to regularize or normalize the predictions</li> </ul> The problem with the latter approach is that position bias is dynamic in real world systems, making it cumbersome to separately estimate it. Hence this paper adopts the former approach that is efficient and adapts to training data distribution changes as we train the main model. This paper treats position bias as a bias term that is added to the prediction logit for each task k, before the logit is fed into a final <code>sigmoid</code> to generate a click probability. The position bias term is generated by passing the <code>position</code> feature and other features like <code>device info</code> into a shallow tower. Device info is included because different position bias is observed on different types of devices. Note: The paper is not clear on how the position and device are represented as features. Naively, I imagine that a small embedding (say of <code>dim 32</code>) is used to encode each position and each device value. Since the position bias is strongest at the top positions, we could cap the number of position embeddings to say <code>50</code>, and assign all later positions to the final position embedding. Then, we can concatenate the position and device embedding to form an embedding of say <code>dim 64</code>, and feed that into a shallow ReLU layer to generate the position bias logit. Another way is to do element-wise multiplication of the position and device embedding. <h2 id="training"><a class="header" href="#training">Training</a></h2> The paper does not provide much detail on how exactly training is performed, but they mentioned that point-wise loss was chosen for simplicity. I would imagine the process to be something like: <ul> <li>Start with positive training examples where some subset of engagement or satisfaction metric were observed (e.g. user watched)</li> <li>For each positive training example, sample one or more random negative(s) from items that were shown to the user in that context but user did not click</li> <li>Train the model to predict the multi-task outcome for each of these examples in a point-wise manner</li> </ul> <h2 id="experiments-2"><a class="header" href="#experiments-2">Experiments</a></h2> Online experiments were performed on YouTube. AUC was chosen for classification tasks and squared error for regression tasks. <ul> <li>The paper found that using the shallow tower for position bias increases AUC by <code>0.24%</code></li> <li>The Mixture of Experts approach is better than the shared bottom approach for the same number of model parameters.</li> </ul> <h2 id="takeaways"><a class="header" href="#takeaways">Takeaways</a></h2> This paper is scant on details on the actual ranker and features used, but makes a good case for the two proposed components of (i) mixture of experts and (ii) shallow tower for position bias. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="lee-2020---large-scale-video-representation-learning"><a class="header" href="#lee-2020---large-scale-video-representation-learning">Lee 2020 - Large Scale Video Representation Learning</a></h1> <a href="https://openaccess.thecvf.com/content_CVPR_2020/papers/Lee_Large_Scale_Video_Representation_Learning_via_Relational_Graph_Clustering_CVPR_2020_paper.pdf">Large Scale Video Representation Learning via Relational Graph Clustering</a> This paper proposes an efficient method to learn video representations from relational graph data (i.e. user item interactions). Specifically, it learns a small-ish embedding model that transforms raw audio-visual features for a given video into a vector representation that is useful for downstream recommendation tasks. This method seems to be in use at YouTube at least until 2023 as it was mentioned in <a href="papers/./singh_2023.html">Singh 2023 - Better Generalization with Semantic IDs</a>. This paper is in the genre of metric learning of similarity metric between videos. Similar to <a href="papers/./schroff_2015.html">FaceNet</a>, it uses triplet contrastive loss to push related videos together and random videos apart. The contribution of the paper is a hierarchical clustering approach to sample smart negatives which they show to be much more informative for learning than random negatives. <h2 id="setup-10"><a class="header" href="#setup-10">Setup</a></h2> We start with a raw input representation of a given video x∈Rd, where x could be a concatenation of various raw input features or a representation from some off-the-shelf pretrained embedding model. We are also given a relational graph G=(V,E) where each node is a video and each edge (x1,x2)∈E represents some relationship between two videos. The edge weight can be binary or real numbers. We may obtain such edge relationships from implicit feedback, e.g. how frequently two videos are co-watched, co-clicked, co-searched, etc. The aim is to learn a representation z∈Rk where k is much smaller than d, such that z1Tz2∼1 if they are related and 0 otherwise. <h2 id="method-1"><a class="header" href="#method-1">Method</a></h2> The relational graph is pre-processed with a hierarchical clustering algorithm. A hierarchical algorithm is chosen so that we can sample negatives with varying levels of difficulty later on. The paper uses <a href="papers/./bateni_2017.html">Affinity Clustering</a>, although it notes that any suitable clustering algorithm works as well. At a high level, affinity clustering at each step chooses the lowest outgoing edge-weight to add from each cluster. Some desirable properties of affinity clustering are (i) tends to produce clusters of similar size, which helps negative sampling to be consistent across clusters and (ii) easily parallelizable. Once we have the relational graph G and clustering C, the training proceeds as follows: <ul> <li>Construct triplets: <ul> <li>Sample a random <code>anchor</code> video from all videos</li> <li>Choose a <code>positive</code> video from its neighbours in G</li> <li>Sample a <code>negative</code> video at a desired level from sibling clusters in C</li> </ul> </li> <li>Compute the distances for each triplet i: <ul> <li>Anchor-positive distance: dap:=∣∣f(xia)−f(xip)∣∣2</li> <li>Anchor-negative distance: dan:=∣∣f(xia)−f(xin)∣∣2</li> </ul> </li> <li>Perform online semi-hard negative mining: <ul> <li>In each mini-batch, resample negatives to choose the hardest semi-hard negative in the mini-batch for each row</li> <li>Recall that a semi-hard negative is where 0<dan−dap≤α</li> <li>So we are choosing negatives such that dan is as close as possible to dap without being lower</li> <li>Note that this means we need to compute dan for all anchor-negative pairs in the mini-batch</li> </ul> </li> <li>Optimize for the following objective: mini=1∑N[dapi−dani+α]+</li> </ul> This is essentially the same procedure as <a href="papers/./schroff_2015.html">FaceNET</a>, except the smart and dynamic choice of negatives. Also note that since we only train with semi-hard negatives, the inner term 0≤dap−dan+α<α, so the + operator on the loss can actually be ignored. Denote the anchor video as <code>A</code>. The smart negative sampling from sibling clusters is defined as: <ul> <li>At <code>L0</code>, negatives are chosen from all the descendants of P(A) from the affinity tree</li> <li>At <code>L1</code>, negatives are chosen from all the descendants of Gp(A):=P(P(A))</li> <li>And so on</li> </ul> <h2 id="experiments-3"><a class="header" href="#experiments-3">Experiments</a></h2> The embedding network is simply a fully connected MLP with layers of dimension <code>4,000</code> and <code>256</code> respectively. Each video is preprocessed into an embedding representing raw audio-visual features using techniques like fourier transform, feeding into pre-trained audio-visual ResNets etc. and fed into the MLP to get a resulting embedding. The relational graph is constructed by adding an edge between a pair of videos if they are frequently co-watched by multiple users. Specifics are not provided but I presume the rule is something like <code>co-watched by at least n user pairs</code>. In other words, the edge rule sounds fairly stringent to minimize the number of false edges. The videos are then split into training and test with a <code>7:3</code> split. Each mini-batch comprises <code>9,600</code> triplets and the margin is set to <code>0.5</code>. The learning rate starts at <code>0.1</code> and decays by <code>0.98</code> for every <code>300k</code> steps. The evaluation simulates a cold start scenario, where both the query video and candidate videos are all taken from the unseen test set. For each query video, retrieval is performed using the model (using cosine similarity) on all videos from the test set. The NDCG and MAP are then computed based on whether we successfully retrieved relevant videos for the query video based on the true relational graph G. Findings: <ul> <li>Smart negatives are significantly better than random negatives. To further prove this point, the authors tracked the % of negatives that remain the same after the hard negative re-sampling step, and showed that a much larger % of smart negatives were retained (10% to 15%) than random negatives (close to 1 / <code>batch_size</code>). Interestingly, the usage % of smart negatives actually increases steadily as training goes on. This could be because at the start, many smart negatives are too difficult for the model and do not qualify as semi-hard. As the model learns, it is able to handle more smart negatives and the usage % increases.</li> <li>Online semi-hard negative mining is important. The authors found that using smart negatives without the online mining step does not perform well. This could be because the smart negatives are too difficult for the model at the beginning and the model fails to learn. The authors suggest that an alternative to the online mining step is curriculum learning, where easier negatives are provided at the start and difficulty is increased gradually.</li> <li>More difficult negatives work better. The authors tried training at <code>L0</code>, <code>L1</code> and <code>L2</code> difficulty levels respectively, and found that training at <code>L0</code> consistently does the best. This kind of defeats the point of the hierarchical clustering, but I suppose it depends on the use case and in other settings other difficulty levels may work better.</li> </ul> <h2 id="takeaways-1"><a class="header" href="#takeaways-1">Takeaways</a></h2> This paper proposes a scaleable and easy-to-understand method of item representation learning. Obviously, it can be extended to other modalities such as text so long as we find a reasonable way to represent the raw inputs. For text, we can probably directly fine-tune some BERT encoder using the same procedure. Nevertheless, note that this paper alone would not provide optimal performance if we use the embeddings here directly for retrieval. One big reason is that we are considering an undirected relational graph, so we learn only symmetric relationships, whereas in retrieval item sequence often matters. For example, people often progress from beginner videos to more advanced videos for a particular subject, and this paper would not be able to capture such relationships. Hence the embeddings in this paper have to be fed into another retrieval or ranking model for optimal performance. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="he-2020---lightgcn"><a class="header" href="#he-2020---lightgcn">He 2020 - LightGCN</a></h1> <a href="https://arxiv.org/pdf/2002.02126.pdf">He 2020 - LightGCN</a> is a simple and effective Graph Convolution Network for recommendation. LightGCN is an adaptation of Graph Convolutional Neural Networks (GCN) to the task of recommendations. In a typical Convolutional Neural Network for vision, convolution aggregations (such as linear projections, pooling, average) are applied to a neighbourhood of pixels that are near to one another. The aggregations transform the raw pixel values into a hidden layer of "embedding" values, and the next layer of aggregations is applied to the hidden layer, allowing the CNN to learn more abstract features with each increasing layer. A GCN uses essentially the same idea, except that the definition of neighbourhood of a node A are the neighbouring nodes that are connected by an edge to A. The GCN thus allows us to train node embeddings on all types of graphical data, such as social networks, user-item interactions etc. <h2 id="setting"><a class="header" href="#setting">Setting</a></h2> This paper tackles the task of collaborative filtering without features, i.e. making recommendations purely from the user and item <code>id</code>. Also, no negative samples are required - all we need is edges between users and items based on some form of implicit interaction. <h2 id="neural-graph-collaborative-filtering-ngcf"><a class="header" href="#neural-graph-collaborative-filtering-ngcf">Neural Graph Collaborative Filtering (NGCF)</a></h2> The LightGCN model is essentially a simplification of the NGCF model, so the paper starts here. Btw, there are some overlaps between LightGCN authors and NGCF authors. The setup is as follows: <ul> <li>Each user and item are embedded from their <code>id</code> -> <code>embedding</code></li> <li>Let eu(0) denote the ID embedding of user u and ei(0) denote the ID embedding of item i</li> </ul> NGCF uses the user-item interaction graph (derived from data) to propagate the embeddings as follows: eu(k+1)ei(k+1)=σ(W1(k)eu(k)+i∈Nu∑∣Nu∣∣Ni∣<svg xmlns="http://www.w3.org/2000/svg" width='400em' height='1.28em' viewBox='0 0 400000 1296' preserveAspectRatio='xMinYMin slice'><path d='M263,681c0.7,0,18,39.7,52,119 c34,79.3,68.167,158.7,102.5,238c34.3,79.3,51.8,119.3,52.5,120 c340,-704.7,510.7,-1060.3,512,-1067 l0 -0 c4.7,-7.3,11,-11,19,-11 H40000v40H1012.3 s-271.3,567,-271.3,567c-38.7,80.7,-84,175,-136,283c-52,108,-89.167,185.3,-111.5,232 c-22.3,46.7,-33.8,70.3,-34.5,71c-4.7,4.7,-12.3,7,-23,7s-12,-1,-12,-1 s-109,-253,-109,-253c-72.7,-168,-109.3,-252,-110,-252c-10.7,8,-22,16.7,-34,26 c-22,17.3,-33.3,26,-34,26s-26,-26,-26,-26s76,-59,76,-59s76,-60,76,-60z M1001 80h400000v40h-400000z'/></svg>1(W1(k)ei(k)+W2(k)(ei(k)⊙eu(k))))=σ(W1(k)ei(k)+u∈Ni∑∣Nu∣∣Ni∣<svg xmlns="http://www.w3.org/2000/svg" width='400em' height='1.28em' viewBox='0 0 400000 1296' preserveAspectRatio='xMinYMin slice'><path d='M263,681c0.7,0,18,39.7,52,119 c34,79.3,68.167,158.7,102.5,238c34.3,79.3,51.8,119.3,52.5,120 c340,-704.7,510.7,-1060.3,512,-1067 l0 -0 c4.7,-7.3,11,-11,19,-11 H40000v40H1012.3 s-271.3,567,-271.3,567c-38.7,80.7,-84,175,-136,283c-52,108,-89.167,185.3,-111.5,232 c-22.3,46.7,-33.8,70.3,-34.5,71c-4.7,4.7,-12.3,7,-23,7s-12,-1,-12,-1 s-109,-253,-109,-253c-72.7,-168,-109.3,-252,-110,-252c-10.7,8,-22,16.7,-34,26 c-22,17.3,-33.3,26,-34,26s-26,-26,-26,-26s76,-59,76,-59s76,-60,76,-60z M1001 80h400000v40h-400000z'/></svg>1(W1(k)eu(k)+W2(k)(ei(k)⊙eu(k)))) Some notes about the propagation equations above: <ul> <li>eu(k+1) and ei(k+1) denote the embedding of user u and item i respectively after <code>k</code> layers of propagation</li> <li>σ is a non-linear activation function</li> <li>Nu denotes the set of items that interacted with user u. For instance, it could be all the items the user purchased within the past 3 months. Ni is the set of users defined in a similar way.</li> <li>W1(k) and W2(k) are trainable weights</li> </ul> Intuitively for a given user, the equation propagates (i) the user embeddings itself (<code>order-1</code>), (ii) the embeddings of neighbouring items (<code>order-1</code>) and (iii) the hadamard interaction between the user and the neighbouring items (<code>order-2</code>). And likewise for the item embeddings. Note that <no neighbourhood sampling> is performed - the entire neighbour set is taken per node. Finally, after training the network of L layers, we obtain L+1 embeddings for each user and item. The embeddings are concatenated as such eu=[eu(0),..., eu(L)] and ei=[ei(0),..., ei(L)] where eu,ei are vectors of dimension RkL. Prediction scores for the match between user u and item i are then computed via the inner product ⟨eu,ei⟩. <h2 id="problem-with-ngcf"><a class="header" href="#problem-with-ngcf">Problem With NGCF</a></h2> The authors argue that NGCF is unnecessarily complicated because traditionally, the base embedding layer eu(0),ei(0) is derived from rich semantic features such as embedding the title of papers etc. This justifies the usage of the activation function σ and the projection weights W1(k),W2(k) etc. to learn a transformation of the semantic features. In contrast, for the collaborative filtering setting, the embeddings are arbitrary numbers tied to each user or item ID. Hence, performing multiple non-linear transformations will not lead to better feature learning. <Note>: I'm not fully convinced by this argument, although the empirical results do support it. I agree with the argument to the extent that the base embedding layer is arbitrary, but imo NGCF can still learn a bigger representation space of models through its non-linear transformations. The problem seems to be more that (i) the richer feature representation is not very useful and (ii) the additional complexity makes the model harder to learn. <h2 id="lightgcn-forward-propagation"><a class="header" href="#lightgcn-forward-propagation">LightGCN Forward Propagation</a></h2> In LightGCN, we essentially remove the non-linear activation and weight projections. The propagation equations simplify to the following: eu(k+1)ei(k+1)=i∈Nu∑∣Nu∣∣Ni∣<svg xmlns="http://www.w3.org/2000/svg" width='400em' height='1.28em' viewBox='0 0 400000 1296' preserveAspectRatio='xMinYMin slice'><path d='M263,681c0.7,0,18,39.7,52,119 c34,79.3,68.167,158.7,102.5,238c34.3,79.3,51.8,119.3,52.5,120 c340,-704.7,510.7,-1060.3,512,-1067 l0 -0 c4.7,-7.3,11,-11,19,-11 H40000v40H1012.3 s-271.3,567,-271.3,567c-38.7,80.7,-84,175,-136,283c-52,108,-89.167,185.3,-111.5,232 c-22.3,46.7,-33.8,70.3,-34.5,71c-4.7,4.7,-12.3,7,-23,7s-12,-1,-12,-1 s-109,-253,-109,-253c-72.7,-168,-109.3,-252,-110,-252c-10.7,8,-22,16.7,-34,26 c-22,17.3,-33.3,26,-34,26s-26,-26,-26,-26s76,-59,76,-59s76,-60,76,-60z M1001 80h400000v40h-400000z'/></svg>ei(k)=u∈Ni∑∣Nu∣∣Ni∣<svg xmlns="http://www.w3.org/2000/svg" width='400em' height='1.28em' viewBox='0 0 400000 1296' preserveAspectRatio='xMinYMin slice'><path d='M263,681c0.7,0,18,39.7,52,119 c34,79.3,68.167,158.7,102.5,238c34.3,79.3,51.8,119.3,52.5,120 c340,-704.7,510.7,-1060.3,512,-1067 l0 -0 c4.7,-7.3,11,-11,19,-11 H40000v40H1012.3 s-271.3,567,-271.3,567c-38.7,80.7,-84,175,-136,283c-52,108,-89.167,185.3,-111.5,232 c-22.3,46.7,-33.8,70.3,-34.5,71c-4.7,4.7,-12.3,7,-23,7s-12,-1,-12,-1 s-109,-253,-109,-253c-72.7,-168,-109.3,-252,-110,-252c-10.7,8,-22,16.7,-34,26 c-22,17.3,-33.3,26,-34,26s-26,-26,-26,-26s76,-59,76,-59s76,-60,76,-60z M1001 80h400000v40h-400000z'/></svg>eu(k)=∣Nu∣<svg xmlns="http://www.w3.org/2000/svg" width='400em' height='1.28em' viewBox='0 0 400000 1296' preserveAspectRatio='xMinYMin slice'><path d='M263,681c0.7,0,18,39.7,52,119 c34,79.3,68.167,158.7,102.5,238c34.3,79.3,51.8,119.3,52.5,120 c340,-704.7,510.7,-1060.3,512,-1067 l0 -0 c4.7,-7.3,11,-11,19,-11 H40000v40H1012.3 s-271.3,567,-271.3,567c-38.7,80.7,-84,175,-136,283c-52,108,-89.167,185.3,-111.5,232 c-22.3,46.7,-33.8,70.3,-34.5,71c-4.7,4.7,-12.3,7,-23,7s-12,-1,-12,-1 s-109,-253,-109,-253c-72.7,-168,-109.3,-252,-110,-252c-10.7,8,-22,16.7,-34,26 c-22,17.3,-33.3,26,-34,26s-26,-26,-26,-26s76,-59,76,-59s76,-60,76,-60z M1001 80h400000v40h-400000z'/></svg>1i∈Nu∑∣Ni∣<svg xmlns="http://www.w3.org/2000/svg" width='400em' height='1.28em' viewBox='0 0 400000 1296' preserveAspectRatio='xMinYMin slice'><path d='M263,681c0.7,0,18,39.7,52,119 c34,79.3,68.167,158.7,102.5,238c34.3,79.3,51.8,119.3,52.5,120 c340,-704.7,510.7,-1060.3,512,-1067 l0 -0 c4.7,-7.3,11,-11,19,-11 H40000v40H1012.3 s-271.3,567,-271.3,567c-38.7,80.7,-84,175,-136,283c-52,108,-89.167,185.3,-111.5,232 c-22.3,46.7,-33.8,70.3,-34.5,71c-4.7,4.7,-12.3,7,-23,7s-12,-1,-12,-1 s-109,-253,-109,-253c-72.7,-168,-109.3,-252,-110,-252c-10.7,8,-22,16.7,-34,26 c-22,17.3,-33.3,26,-34,26s-26,-26,-26,-26s76,-59,76,-59s76,-60,76,-60z M1001 80h400000v40h-400000z'/></svg>ei(k)=∣Ni∣<svg xmlns="http://www.w3.org/2000/svg" width='400em' height='1.28em' viewBox='0 0 400000 1296' preserveAspectRatio='xMinYMin slice'><path d='M263,681c0.7,0,18,39.7,52,119 c34,79.3,68.167,158.7,102.5,238c34.3,79.3,51.8,119.3,52.5,120 c340,-704.7,510.7,-1060.3,512,-1067 l0 -0 c4.7,-7.3,11,-11,19,-11 H40000v40H1012.3 s-271.3,567,-271.3,567c-38.7,80.7,-84,175,-136,283c-52,108,-89.167,185.3,-111.5,232 c-22.3,46.7,-33.8,70.3,-34.5,71c-4.7,4.7,-12.3,7,-23,7s-12,-1,-12,-1 s-109,-253,-109,-253c-72.7,-168,-109.3,-252,-110,-252c-10.7,8,-22,16.7,-34,26 c-22,17.3,-33.3,26,-34,26s-26,-26,-26,-26s76,-59,76,-59s76,-60,76,-60z M1001 80h400000v40h-400000z'/></svg>1u∈Ni∑∣Nu∣<svg xmlns="http://www.w3.org/2000/svg" width='400em' height='1.28em' viewBox='0 0 400000 1296' preserveAspectRatio='xMinYMin slice'><path d='M263,681c0.7,0,18,39.7,52,119 c34,79.3,68.167,158.7,102.5,238c34.3,79.3,51.8,119.3,52.5,120 c340,-704.7,510.7,-1060.3,512,-1067 l0 -0 c4.7,-7.3,11,-11,19,-11 H40000v40H1012.3 s-271.3,567,-271.3,567c-38.7,80.7,-84,175,-136,283c-52,108,-89.167,185.3,-111.5,232 c-22.3,46.7,-33.8,70.3,-34.5,71c-4.7,4.7,-12.3,7,-23,7s-12,-1,-12,-1 s-109,-253,-109,-253c-72.7,-168,-109.3,-252,-110,-252c-10.7,8,-22,16.7,-34,26 c-22,17.3,-33.3,26,-34,26s-26,-26,-26,-26s76,-59,76,-59s76,-60,76,-60z M1001 80h400000v40h-400000z'/></svg>eu(k) The final representation of each node v (whether user or item) is then a weighted sum of its hidden representation across all layers: ev=k=0∑Kαk⋅ev(k) Although αk could be a parameter to be optimized, the authors propose just setting αk=1/(K+1) for simplicity. Noticeably, the forward propagation does not include the self-connection from the previous layer, i.e. the update for euk+1 does not explicitly include euk, which other papers like <a href="papers/./hamilton_2017.html">GraphSAGE</a> argue is important. The authors argue that because they use a weighted sum of hidden representations across all layers, this essentially is equivalent to including self-connections, so that is no longer necessary. <h2 id="loss-2"><a class="header" href="#loss-2">Loss</a></h2> The only trainable parameters of the model are the embeddings at the 0th layer, i.e. E(0). The authors propose using <a href="papers/rendle_2009.html">Bayesian Personalized Ranking</a> loss, which is a pairwise loss that encourages the score of a neighbour to be higher than the score of an unobserved, randomly sampled counterpart. LBPR=−u=1∑Mi∈Nu∑j∈/Nu∑ln σ(y^ui−y^uj)+λ∣∣E(0)∣∣2 In contrast to NGCF and other GCN approaches, the authors do not use dropout as a regularizer. Instead, they think the L2 regularization on the embedding layer is sufficient, as these are the only parameters in the model. Training of the model is done in a mini-batch manner, where batches of <code>(user, item)</code> tuples are drawn, negative items sampled, and the loss evaluated. <h2 id="ablation-studies"><a class="header" href="#ablation-studies">Ablation Studies</a></h2> The paper has a few ablation findings: <ol> <li>Symmetric Normalization is important, i.e. it is important in the forward propagation to divide by ∣Nu∣∣Ni∣<svg xmlns="http://www.w3.org/2000/svg" width='400em' height='1.28em' viewBox='0 0 400000 1296' preserveAspectRatio='xMinYMin slice'><path d='M263,681c0.7,0,18,39.7,52,119 c34,79.3,68.167,158.7,102.5,238c34.3,79.3,51.8,119.3,52.5,120 c340,-704.7,510.7,-1060.3,512,-1067 l0 -0 c4.7,-7.3,11,-11,19,-11 H40000v40H1012.3 s-271.3,567,-271.3,567c-38.7,80.7,-84,175,-136,283c-52,108,-89.167,185.3,-111.5,232 c-22.3,46.7,-33.8,70.3,-34.5,71c-4.7,4.7,-12.3,7,-23,7s-12,-1,-12,-1 s-109,-253,-109,-253c-72.7,-168,-109.3,-252,-110,-252c-10.7,8,-22,16.7,-34,26 c-22,17.3,-33.3,26,-34,26s-26,-26,-26,-26s76,-59,76,-59s76,-60,76,-60z M1001 80h400000v40h-400000z'/></svg>. Omitting either one leads to performance degradation. Note that in <a href="papers/./hamilton_2017.html">GraphSAGE</a>, the <code>GraphSAGE-mean</code> variant essentially does ∑i∈Nuei(k)/∣Nu∣, i.e. it only normalizes by the user degree. I suppose normalizing by the item degree as well penalizes popular items, so it could be useful.</li> <li>Layer combination is important for robustness as we increase the number of layers, i.e. instead of just taking eK as the final embeddings, it is useful to take the element-wise mean of the embeddings at each layer. This might be analogous to the impact of including self connections.</li> </ol> <h2 id="cornac-implementation-1"><a class="header" href="#cornac-implementation-1">Cornac Implementation</a></h2> Cornac has a torch implementation of LightGCN: <ul> <li><a href="https://github.com/PreferredAI/cornac/blob/master/cornac/models/lightgcn/recom_lightgcn.py">Training</a></li> <li><a href="https://github.com/PreferredAI/cornac/blob/master/cornac/models/lightgcn/lightgcn.py">Model / Graph Sampling</a></li> </ul> The code relies on the <a href="https://github.com/dmlc/dgl">dmlc/dgl package</a> for constructing the bipartite user-item graph which will be used to compute neighbourhoods. The <code>construct_graph</code> function works as follows: <ul> <li><code>user_indices</code> and <code>item_indices</code> are lists of the same length where each element at index <code>i</code> contains a pair of user, item that interacted</li> <li>A <code>dgl.heterograph</code> is constructed with both directions: <ul> <li><code>("user", "user_item", "item")</code> represents user -> item direction</li> <li><code>("item", "item_user", "user")</code> represents item -> user direction</li> <li>Hence there are two node types and two edge types in the graph</li> </ul> </li> <li>Starting with the <code>user->item</code> direction: <ul> <li><code>src</code> and <code>dst</code> are torch tensors containing the users and items respectively that interacted, both of length <code>M</code></li> <li><code>dst_degree</code> is a torch float tensor of length <code>M</code> containing the number of users interacting with each item in <code>dst</code></li> <li><code>src_degree</code> is a torch float tensor of length <code>M</code> containing the number of items interacting with each user in <code>src</code></li> </ul> </li> </ul> At model initialization, <code>self.feature_dict</code> is initialized with xavier initialization as follows. Note that because we have a heterograph, the nodes are defined as a dictionary of the form <code>node_type: feature_tensor</code>. <pre><code class="language-python"> self.feature_dict = { "user": user_embed, # (n_users, embed_dim) "item": item_embed, # (n_items, embed_dim) } </code></pre> The <code>GCNLayer</code> class represents one layer of the message passing network. <h2 id="references-4"><a class="header" href="#references-4">References</a></h2> <ul> <li><a href="https://medium.com/stanford-cs224w/lightgcn-with-pytorch-geometric-91bab836471e">Medium Post by Hikaru Hotta and Ada Zhou</a></li> <li><a href="https://arxiv.org/pdf/2002.02126.pdf">LightGCN Paper</a></li> <li><a href="https://arxiv.org/pdf/1905.08108.pdf">NGCF Paper</a></li> </ul> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="lewis-2020---retrieval-augmented-generation"><a class="header" href="#lewis-2020---retrieval-augmented-generation">Lewis 2020 - Retrieval Augmented Generation</a></h1> <a href="https://ar5iv.labs.arxiv.org/html/2005.11401">Lewis 2020 - Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks</a> This paper proposes a way to fine-tune a parametric seq2seq transformers (e.g. GPT) with a non-parametric memory through dense retrieval. The main idea is to extend parametric memory (i.e. the "knowledge" that is stored within the LLM floating point parameters) of the seq2seq model by coupling it with retrieval of documents from a vector database (dubbed non-parametric memory, using Wikipedia articles in the original paper). We can update the non-parametric database's knowledge as the world changes. The paper argues that this setup is ideal for knowledge-intensive tasks such as open question answering, fact verification etc., where it is impossible to store all knowledge in parametric memory. <h2 id="setup-11"><a class="header" href="#setup-11">Setup</a></h2> Given an input sequence of tokens x, we wish to retrieve contextual documents z and use them as additional context when generating target sequence y. We have two models: <ul> <li>Retriever pη(z∣x) which returns top-k truncated probability distributions over text passages. Truncated probably means that the probability is normalized such that the top-k probabilities sum up to <code>1</code> for each retrieval.</li> <li>Generator pθ(yi∣x,z,y1:i−1) parametrized by θ which generates tokens in an auto-regressive manner.</li> </ul> The goal is to train both models end-to-end on a fine-tuning corpus of input sequence / output sequence pairs. The loss objective is to minimize the negative log-likelihood −log p(yj∣xj) . <h2 id="training-1"><a class="header" href="#training-1">Training</a></h2> The paper proposes two different models to achieve the end-to-end prediction. <ol> <li>RAG-Sequence. In this setting, we retrieve <code>k</code> documents and keep using them to generate the entire target sequence.</li> </ol> pRAG−Sequence(y∣x)∼z ∈ topk{pη(⋅∣x)}∑pη(z∣x)pθ(y∣x,z)=z ∈ topk{pη(⋅∣x)}∑pη(z∣x)i∏Npθ(yi∣x,z,y1:i−1) <ul> <li>Note that we are marginalizing (or taking a weighted combination) over the truncated distribution of z, implying that we trust each document according to its retrieval probability in the final probability for generating each token.</li> <li>The expression pθ(yi∣x,z,y1:i−1) just means that we generate the target sequence with an input sequence that is a concatenation of x, z and y1:i−1.</li> <li>The retrieval is done using a BERT encoder using Maximum Inner Product Search (MIPS). To avoid re-indexing, the document vectors are held constant whilst the query encoder is trained in the above end-to-end fashion.</li> <li>There is no explicit supervision on what documents are to be retrieved. Intuitively, if a document is useful for generating the correct tokens, the loss objective would encourage pη(z∣x) to be larger, thus encouraging the retriever to retrieve more relevant documents.</li> <li>Another interesting way to think about this setup: suppose the generator just returns the retrieved document (token for token) and <code>k=1</code>, and the input / output pairs are <code>anchor-positive</code> pairs in a standard retrieval setting. Then we can see that this matches the standard retrieval training objectives but without negative sampling. Thus it seems that the token prediction task is sufficient to generate negative signals for non-useful documents such that explicit negatives are not needed.</li> </ul> <ol start="2"> <li>RAG-token. In contrast, RAG-token retrieves <code>k</code> documents at each time step, allowing us to sample new documents according to the needs at each decoding step.</li> </ol> pRAG−Token(y∣x)∼i∏Nzi ∈ topk{pη(⋅∣x,y1:i−1)}∑pη(z∣x)pθ(yi∣x,zi,y1:i−1) Note that the retrieved context is now zi which varies at each time step. The change in retrieval at each step seems to add complexity during training. <h2 id="ablation-1"><a class="header" href="#ablation-1">Ablation</a></h2> <ul> <li>Increasing number of retrieved documents. <code>5</code> or <code>10</code> documents are used for retrieval. Ablation shows that performance increases monotonically (albeit diminishingly) for <code>RAG-sequence</code> with increasing number of retrieved documents.</li> <li>Learned Retrieval is useful. The authors try freezing the retriever and compare it against the setting of allowing the retriever to learn. They find that generally learned retrieval improves results significantly.</li> <li>RAG generates more diverse outputs. They measure the ratio of <code>distinct tri-grams / total tri-grams</code> and find that RAG-decoding compared to normal decoding is more diverse.</li> </ul> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="wang-2020---dcnv2"><a class="header" href="#wang-2020---dcnv2">Wang 2020 - DCNv2</a></h1> <a href="https://arxiv.org/abs/2008.13535">DCN V2: Improved Deep & Cross Network and Practical Lessons for Web-scale Learning to Rank Systems</a> This paper proposes a neural network architecture called Deep and Cross Networks for effectively learning feature crosses compared to a standard feedforward MLP. <h2 id="background-2"><a class="header" href="#background-2">Background</a></h2> Effective feature crossing is essential in many ML tasks, especially for search and recommendation. For example, a country feature crossed with language is more informative than either alone. Manually searching for good feature crosses is a combinatorial exercise which is intensive. Deep neural networks in the form of MLPs are generally viewed as universal function approximators in the limiting case. But in the finite depth, they are often incapable of fully modelling feature crosses effectively (say using a simulated dataset). Traditionally, Factorization Machines seek to overcome the feature combination issue by embedding a sparse feature with large dimensionality into a small dense vector. Feature crosses between features is then performed using a dot product between features. However, we are limited in expressiveness to order-2 crosses and the number of feature crosses can still be large with a large number of features (after densification). The other limitation of FMs is that we require the embedding dimension of each feature be the same, which is limiting if our features have different cardinality needs. Traditional approach is to mix implicit crossing (i.e. fitting a DNN to the features) with explicit crossing (say with the FM approach, or by multiplying raw features together etc.). The implicit network and explicit network are usually in parallel and the output is added together to form the final prediction. <h2 id="dcnv1"><a class="header" href="#dcnv1">DCNv1</a></h2> The same authors proposed DCNv1 in 2017, and it is useful to see how it evolved. The method is as follows: <ul> <li>Let x[0] represent a concatenated feature vector at layer 0. That is, we lay out each dense or sparse feature sideways and concatenate them into a long feature vector. Let x[0]∈Rd.</li> <li>We have a cross and deep layer in parallel <ul> <li>The deep side is simply a standard MLP, i.e. h1=Relu(w[0]x[0]+b[0])</li> <li>The cross side is where the magic happens: x[l+1]=x[0]x[l]Tw[l]+b[l]+x[l]</li> <li>Note that w,b∈Rd are learnable parameters</li> <li>x[0]⋅x[l]T is a d×d matrix of rank <code>1</code>. At layer <code>1</code>, it comprises all the pairwise crosses, e.g. x1[0]⋅x1[l], x1[0]⋅x2[l], ...</li> <li>Hence the transformation at each cross layer is rank <code>1</code>.</li> </ul> </li> <li>As we increase the number of cross layers, we will get feature crosses of increasing polynomial degree. We end up with a d-dimensional feature vector x[l]∈Rd which comprises complex weighted polynomial feature crosses of polynomial degrees 1,...,l.</li> <li>At the final layer, we have xwide[L],xdeep[L]∈Rd. The two features are concatenated together to form a vector R2d, which can then be fitted to a classifier head for final predictions.</li> </ul> <h2 id="dcnv2"><a class="header" href="#dcnv2">DCNv2</a></h2> The criticism of DCNv1 is that the transformation at each cross layer is rank <code>1</code> and hence not expressive enough. DCNv2 tries to make the cross layer more expressive while still making it parameter-efficient. The cross layer formulation of DCNv2 is: x[l+1]=x[0]⊙(W[l]x[l]+b[l])+x[l] Where W[l]∈Rd×d,b[l]∈Rd. To see how it compares to DCNv1, we can let W[l]=u[l]⋅v[l]T be a rank <code>1</code> matrix and b[l]=0. Furthermore, we set u[l]=1. Then we have: x[l+1]=x[0]⊙(u[l]⋅v[l]Tx[l])+x[l]=(v[l]Tx[l])(u[l]⊙x[0])+x[l]=(v[l]Tx[l])x[0]+x[l] Note that in line 2, we use the fact that v[l]Tx[l] is a scalar and move it out to the left. In line 3 since u[l]=1 we can remove it. Similarly for DCNv1, we can pull out the scalar: x[l+1]=x[0]x[l]Tw[l]+b[l]+x[l]=(x[l]Tw[l])x[0]+b[l]+x[l] We thus see that DCNv2 ends up in exactly the same form as DCNv1 (with just a missing b[l] term). This reformulation helps us see that DCNv1 is DCNv2 when W[l] is rank <code>1</code>. Hence when we allow W[l] to be higher rank, we get more expressiveness than DCNv1. <h2 id="stack-vs-parallel"><a class="header" href="#stack-vs-parallel">Stack vs Parallel</a></h2> In addition to the parallel structure proposed in DCNv1, where a deep MLP runs parallel to the cross network and the final vector is concatenated together, DCNv2 proposes an alternative stacked architecture. In this formulation, we run through the cross layers first to get x[l]∈Rd. Then, this is fed into a deep MLP. The paper says that which architecture performs better depends on the task. <h2 id="loss-3"><a class="header" href="#loss-3">Loss</a></h2> Finally, the loss is computed as standard binary cross entropy wrt the binary labels: L=−N1i=1∑Nyilog(y^i)+(1−yi)log(1−y^i)+λl∑∣∣Wl∣∣22 <h2 id="modifications"><a class="header" href="#modifications">Modifications</a></h2> Using ranking models in production settings usually has strict latency requirements, hence it is important to reduce cost while maintaining accuracy. The paper thus proposes 3 modifications to make the model more efficient. <h3 id="modification-1-low-rank-approximation"><a class="header" href="#modification-1-low-rank-approximation">Modification 1: Low Rank Approximation</a></h3> In practice, the weight matrix W[l] is usually effectively low rank, so it is well motivated to approximate it with smaller matrices U[l]V[l]T, where both U[l],V[l]∈Rd×r,r<<d. So we have: x[l+1]=x[0]⊙(U[l](V[l]⊺x[l])+b[l])+x[l] For the experimental setting in the paper, r=64 was the low rank threshold after which they reported diminishing returns for increasing rank. <h3 id="modification-2-mixture-of-experts"><a class="header" href="#modification-2-mixture-of-experts">Modification 2: Mixture of Experts</a></h3> Instead of just having one expert weight for each cross layer, they propose having multiple experts and then combining the expert outputs together using a gating mechanism. This is analogous to multi-headed attention with multiple heads. The idea is that each expert can learn effective feature crosses in a certain subspace. The input-dependent gating mechanism can then select the appropriate experts for a given input. We have: x[l+1]Ei(x[l])=i=1∑KGi(x[l])Ei(x[l])+x[l]=x[0]⊙(Ui[l](Vi[l]⊺x[l])+b[l]) Where: <ul> <li>Gi(⋅):Rd↦R is the gating function which represents the input-dependent weight of expert i. It can be a learned softmax function.</li> <li>Ei is the expert i. It is simply the earlier equation but with separate weights for each expert i.</li> </ul> <h3 id="modification-3-pre-activation-functions"><a class="header" href="#modification-3-pre-activation-functions">Modification 3: Pre Activation Functions</a></h3> With the low rank approximation, we effectively project the features to a low dimension and project it back up. Instead of immediately projecting back from dimension r to d, we can apply non-linear transformations. This allows the function to learn a richer set of representations. Ei(x[l])=x[0]⊙(Ui[l]⋅g(Ci[l]⋅g(Vi[l]⊺x[l]))+b[l]) Here, g(⋅) represents any non-linear activation function (like Relu) and Ci∈Rr×r is a learned weight. In the paper, the sigmoid function was chosen. In practice, the tensorflow implementation seems to incorporate (i) low rank approximation and (ii) pre activation function, but does not do the mixture of experts. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="gao-2021---gradcache"><a class="header" href="#gao-2021---gradcache">Gao 2021 - GradCache</a></h1> <a href="https://arxiv.org/abs/2101.06983">Scaling Deep Contrastive Learning Batch Size under Memory Limited Setup</a> This paper demonstrates a method to perform contrastive learning for two-tower model training with an arbitrarily large batch size at a constant memory cost, at the expense of slightly longer compute time. <h2 id="background-3"><a class="header" href="#background-3">Background</a></h2> The typical contrastive learning setup for learning a two tower retrieval model is explained by <a href="https://arxiv.org/abs/2004.04906">Karpukhin 2020 - Dense Passage Retrieval</a>, in which we have batches of related <code>(anchor, positive)</code> passages. Contrastive learning sets out to maximize the similarity between the anchor and positive passage, and minimize the similarity between the anchor and all other passages in the mini batch. It has been consistently shown that using larger batch sizes is critical to the performance of this training method, as the number of negatives increases, providing more information to the contrastive learning process. However, using large batch sizes is impractical for most researchers, since the memory cost scales linearly with the batch size. The DPR paper used 8x V100 GPUs to process a batch size of 128, which is not attainable for most outfits. Thus the method in this paper is of great practical significance. <h2 id="setup-12"><a class="header" href="#setup-12">Setup</a></h2> We start with two classes or sets of data, S and T. Typically S may represent a set of string queries and T may represent a set of document texts. We want to learn encoders f and g such that, given s∈S and t∈T, the encoded representations f(s) and g(t) are close if related and far apart if not related. Typically, we set up a contrastive loss as follows. Sample a mini batch of anchors S⊂S and corresponding targets T⊂T, where each element si∈S has a corresponding target tri∈T. The rest of the random samples in T will be used as in-batch negatives. We have an InfoNCE loss as follows: L=−∣S∣1si∈S∑log∑tj∈Texp(f(si))Tg(tj)/τexp(f(si))Tg(tri)/τ Let us also denoted the parameters of f as Θ and parameters of g as Λ. <h2 id="analysis-1"><a class="header" href="#analysis-1">Analysis</a></h2> Now we show how the gradient computation and therefore training can be broken down to mitigate the memory bottleneck. Importantly, note that the main bottleneck in such contrastive training is that increasing the batch size scales linearly with the maximum memory requirement of the forward pass of the large BERT model. This is because we encode all texts in the mini-batch simultaneously and run backpropagation. Hence, we want a method that allows us to batch the forward pass within a mini-batch into mini-mini-batches (lets call it a tiny batch) while still allowing us to get the correct backpropagation gradients. Applying the multivariate chain rule to the loss above, we have that: ∂Θ∂L∂Λ∂L=si∈S∑∂f(si)∂L∂Θ∂f(si)=tj∈T∑∂g(tj)∂L∂Λ∂g(tj) From these simple statements, the paper makes two important observations: <ol> <li>The partial derivative ∂Θ∂f(si) only depends on Θ and si. It does not depend on any other anchor or passage. Thus, if we have access to the numerical value of ∂f(si)∂L, we can run backpropagation for ∂Θ∂L independently from all other samples in an arbitrarily small batch.</li> <li>The partial derivative ∂f(si)∂L requires only the numerical values of the encoded representations f(si) for all si∈S and g(tj) for all tj∈T. To compute these values, we don't actually need the computation graph states of the encoder f, we just need the numerical values of all the embeddings.</li> </ol> Note that we can The above statements are focused on f, si and Θ, but similar statements hold for g, tj and Λ. The first statement above allows us to run the expensive gradient updates on a small batch of anchors or passages at a time, which avoids the memory bottleneck of running gradient updates on a large batch size for the large encoders. We can do this so long as we have access to the partial derivatives ∂f(si)∂L. The second statement shows us that computing these partial derivatives is not difficult, because we just need the encoded representations of each anchor and passage. Hence we can batch encode all the anchors and passages in the mini-batch (without gradients), and then use these values to compute this derivative. <h2 id="method-2"><a class="header" href="#method-2">Method</a></h2> The above explanation directly informs the algorithm called GradCache. It works as follows. First, split the large batch into tiny batches which can fit into memory, denoted as S={S1^,S2^,...} and T={T1^,T2^}. Step 1: Graph-less Forward. We run a no-gradient forward pass of each encoder to obtain f(si),g(tj) for all si∈S,tj∈T. We store all the encoded representations in memory. Step 2: Gradient Caching. Using the pre-computed representations in step 1, we run a forward pass to obtain the loss L. We then allow the autograd library to run a backward pass to get gradients ∂f(si)∂L,∂g(tj)∂L for each representation f(si) or g(tj). Note that we are not involving the encoders in this step at all, so memory costs are minimal (just some dot products). <ul> <li>Denote ui=∂f(si)∂L</li> <li>Denote vi=∂g(ti)∂L</li> <li>Store the representation gradient cache {ui,u2,...,v1,v2,...}</li> </ul> Step 3: Tiny Batch Gradient Accumulation. Recall earlier we said that so long as we have access to the partial derivatives from the loss to the embeddings, we can compute gradients for each si or tj in arbitrarily tiny batches. This is the step where we do so. Specifically, we perform gradient accumulation one tiny batch at a time. For the parameters Θ of encoder f: ∂Θ∂L=Sj^∈S∑si∈Sj^∑∂f(si)∂L∂Θ∂f(si)=Sj^∈S∑si∈Sj^∑ui∂Θ∂f(si) The uis are simply looked up from the gradient cache. We perform encoder forward on a tiny batch at each time, multiply with ui and accumulate the gradients. Thus the memory requirement is limited to the encoder forward on a tiny batch, which can be arbitrarily small. Note that the final gradient computed and applied will be equivalent to the original gradients had we directly computed the loss for the large batch. We can see this from the double summation in the equation above, which simply equates to summing over all si∈S. Finally, we perform a similar gradient accumulation for the parameters Λ of encoder g. Once all the sub-gradients are accumulated, the optimizer step is taken to update model parameters as though we had processed the full batch in a single forward backward pass. <h2 id="results-1"><a class="header" href="#results-1">Results</a></h2> The results are based on replicating the Dense Passage Retrieval paper results using a smaller GPU. Note that DPR was trained on 8 GPUs and batch size of 128. <ul> <li>DPR had a top-20 hit rate of <code>78.4</code></li> <li>Using a batch size of 8, which was the largest batch size that fits in memory on an RTX 2080ti, the top-20 hit rate was <code>77.2</code></li> <li>Using gradcache to simulate the batch size of 128, the top-20 hit rate was <code>79.3</code></li> <li>Using gradcache to simulate a larger batch size of 512, the top-20 hit rate was <code>79.9</code></li> </ul> These results demonstrate the importance of a large batch size and the effectiveness of the method to do so on a small GPU. <h2 id="takeaways-2"><a class="header" href="#takeaways-2">Takeaways</a></h2> This paper uses a simple property of the chain rule to separate the computation of the InfoNCE loss gradients into two parts, thereby removing the memory bottleneck. This method is implemented as the <code>CachedMultipleNegativesRankingLoss</code> in <code>SentenceTransformer</code>, and is really useful for those of us without access to large memory GPUs. We can train a model exactly to the performance of an arbitrarily large batch size, at the cost of longer computation time. <h2 id="implementation-3"><a class="header" href="#implementation-3">Implementation</a></h2> The author Luyu Gao has an implementation of GradCache in pytorch. The source code of the main classes are here: <ul> <li><a href="https://github.com/luyug/GradCache/blob/main/src/grad_cache/pytorch_lightning/pl_gradcache.py">PLGradCache</a></li> <li><a href="https://github.com/luyug/GradCache/blob/main/src/grad_cache/grad_cache.py#L16">GradCache</a></li> <li><a href="https://github.com/luyug/GradCache/blob/main/src/grad_cache/context_managers.py">RandContext</a></li> </ul> The main method is in <code>cache_step</code> which computes the loss for a mini batch. We follow the logic below: <ul> <li>Step 1: <code>forward_no_grad</code> is called on the <code>model_inputs</code> to get the encoded representations (or embeddings) of all the input texts: <ul> <li><code>torch.no_grad</code> is used as context manager to avoid gradients</li> <li>In a for loop over the sub-batches: <ul> <li>The model <code>forward</code> method is called on the sub-batch input tensors</li> <li><code>RandContext</code> context manager for this forward pass is initialized and stored in a list of <code>rnd_states</code> <ul> <li>We need to store the random state of both CPUs and GPUs for this forward pass, because we need to exactly replicate the random number state at this point in time later. These random states can affect the behaviour of certain nn layers, especially DropOut.</li> <li>The <code>RandContext</code> object will be used as context manager later on</li> </ul> </li> </ul> </li> <li>The sub-batch representation tensors are concatenated together</li> <li>The <code>model_reps</code> and <code>rnd_states</code> are returned <ul> <li><code>model_reps</code> is appended to a list <code>all_model_reps</code></li> <li><code>rnd_states</code> is appended to a list <code>all_rnd_states</code></li> </ul> </li> </ul> </li> <li>Step 2: <code>build_cache</code> is called to build the cache of gradients. These are the gradients from the loss L to the embeddings f(si),g(tj). <ul> <li><code>compute_loss</code> is called to forward pass from the embeddings to the loss</li> <li><code>backward</code> is called to compute the gradients from the loss to the embeddings</li> <li>For each embedding <code>r</code>, <code>r.grad</code> is accessed to get the gradients</li> <li>The cache is thus <code>[r.grad for r in reps]</code></li> </ul> </li> <li>Step 3: <code>forward_backward</code> is called to accumulate gradients <ul> <li>Firstly, <code>with state</code> is called to restore the random context that we stored earlier. This ensures that the forward pass of the model to get the embedding (this time with gradients) exactly matches the earlier forward pass with no gradients.</li> <li>We obtain the embeddings <code>y</code> with gradients enabled. This corresponds to f(si) in the analysis above.</li> <li>We retrieve the gradients associated with each embedding that we stored earlier in step 2 (call it <code>reps</code>). This corresponds to ui in the analysis above.</li> <li>Now we dot product <code>reps</code> and <code>y</code> to form <code>surrogate</code>. <code>backward</code> is then called on <code>surrogate</code> to get the correct gradients. <ul> <li>This step is a bit tricky, let's look at it in a bit more detail.</li> <li>Recall that our objective is to obtain ∂f(si)∂LΘ∂f(si)</li> <li>Recall that ui is the precomputed numerical value of ∂f(si)∂L</li> <li>Since ui is a constant, it will become a constant multiplier to the gradient on the backward pass</li> <li>Hence by calling backward on <code>surrogate</code>, we will get gradients of the form ui⋅∂Θ∂f(si), which is what we want</li> </ul> </li> <li>After the <code>forward_backward</code> function, the gradients will be accumulated in <code>model.parameters().grad</code>, and the optimizer step can then be taken</li> </ul> </li> </ul> <code>RandContext</code> itself is an interesting context manager to store and load pytorch's internal rng state. <ul> <li>On <code>init</code>, the current cpu and gpu rng states are captured: <ul> <li><code>torch.get_rng_state()</code> gets a byte tensor representing the cpu rng state</li> <li><code>torch.utils.checkpoint.get_device_states(*tensors)</code> looks up the gpu device where the tensors are held, and returns both the <code>gpu_devices</code> and <code>gpu_rng_states</code> across all gpu devices</li> </ul> </li> <li><code>__enter__</code> is triggered when this class is used as context manager, to restore the earlier captured rng state. <ul> <li><code>self._fork = torch.random.fork_rng</code> is called to create a fork of the current torch rng state. This creates an isolated rng environment where we can restore the earlier rng environment without messing up the original rng environment that we entered from.</li> <li><code>self._fork.__enter__()</code> is called to actually enter the forked state</li> <li><code>torch.set_rng_state</code> now sets the cpu rng state to the earlier captured state</li> <li><code>torch.utils.checkpoint.set_device_states</code> similarly sets the gpu rng state to the earlier captured state</li> </ul> </li> <li><code>__exit__</code> is triggered when the context manager is closed. <ul> <li><code>self._fork.__exit__</code> is called to close the isolated rng environment</li> </ul> </li> </ul> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="gao-2021---simple-contrastive-learning-of-sentence-embeddings"><a class="header" href="#gao-2021---simple-contrastive-learning-of-sentence-embeddings">Gao 2021 - Simple Contrastive Learning of Sentence Embeddings</a></h1> <a href="https://ar5iv.labs.arxiv.org/html/2104.08821">Paper Link</a> This paper proposes an unsupervised and supervised approach to fine-tune sentence encoder models to perform semantic textual similarity tasks. The STS tasks refer to a set of tasks taken from SemEval 2012 to 2017 where given two sentences, the model is to output a score between <code>1-5</code> and scored based on correlation to human inputted scores. <h2 id="task-dataset"><a class="header" href="#task-dataset">Task Dataset</a></h2> Example sentence pair with a score of <code>4</code> (taken from <a href="https://aclanthology.org/S16-1081.pdf">SemEval 2016</a>): <ul> <li>In May 2010, the troops attempted to invade Kabul.</li> <li>The US army invaded Kabul on May 7th last year, 2010</li> </ul> Example sentence pair with a score of <code>3</code>: <ul> <li>John said he is considered a witness but not a suspect.</li> <li>"He is not a suspect anymore." John said.</li> </ul> <h2 id="unsupervised-simcse"><a class="header" href="#unsupervised-simcse">Unsupervised SimCSE</a></h2> SimCSE follows the popular framework of contrastive learning using in-batch negatives, where pairs of related sentences are trained to have high similarity whilst having low similarity with other random sentences. The idea of unsupervised SimCSE is simple: given a collection of sentences {xi}i=1m, treat each sentence itself as its own positive pair, and use dropout noise to introduce random perturbations such that the self similarity is not trivially perfect. Specifically, if we denote hiz as the embedding for xi under dropout mask z, then the loss for unsupervised SimCSE for a mini-batch of N sentences is: l1=−log∑j=1Nesim(hjzj,hjzj′)/τesim(hizi,hizi′)/τ Note that τ is the temperature hyperparameter. Importantly, the authors found that setting τ=1 with cosine similarity performs very poorly (<code>64.0</code>), compared to dot product (<code>85.9</code>). However, carefully tuning τ can lead to similar performance (τ=0.05 had a score of <code>86.2</code>). We may view this procedure as data augmentation, analogous to how random pixel distortions and rotations are applied to images to improve computer vision models. The paper shows that this simple unsupervised method significantly outperforms other data augmentation methods. Note that the authors used the default <code>10%</code> dropout for BERT models. <h2 id="supervised-simcse"><a class="header" href="#supervised-simcse">Supervised SimCSE</a></h2> The supervised version follows a similar framework, although the positive pairs are taken from an external dataset. They chose the Natural Language Inference (NLI) datasets, where each example is a triplet of sentences. The <code>premise</code> is denoted xi, the <code>entailment</code> sentence is denoted xi+ and the <code>contradiction</code> sentence is denoted xi−. The loss is then formulated as: l2=−log∑j=1Nesim(hi,hj+)/τ+esim(hi,hj−)/τesim(hi,hi+)/τ <h2 id="ablation-studies-1"><a class="header" href="#ablation-studies-1">Ablation Studies</a></h2> <ul> <li>The paper finds that including <code>contradiction</code> sentences as hard negatives has a small but significant improvement in performance</li> <li>The paper finds that using the <code>[CLS]</code> token or averaging embeddings across the first and last layer does not make much difference</li> </ul> <h2 id="alignment-and-uniformity"><a class="header" href="#alignment-and-uniformity">Alignment and Uniformity</a></h2> <a href="https://ar5iv.labs.arxiv.org/html/2005.10242">Wang and Isola 2020</a> propose two metrics for measuring the effectiveness of an embedding method on a set of documents: <ul> <li>Alignment. Given a distribution of positive pairs of documents ppos, alignment desires the expected distance between embeddings of each pair to be small:</li> </ul> lalign=E(x,x+) ∼ ppos∣∣f(x)−f(x+)∣∣2 <ul> <li>Uniformity. Given any two documents drawn from the corpus, uniformity metric should be small (i.e. distance between them is large).</li> </ul> luniform=log E(x,y) ∼ pdatae−2∣∣f(x)−f(y)∣∣2 A common problem pointed out in training language models is anisotropy, in which embeddings are pushed into a narrow cone in the vector space, which severely limits their expressiveness. The anisotropy problem is naturally connected to uniformity, which aims at distributing embeddings evenly in the space. The authors argue that contrastive learning as proposed in this paper addresses this problem through some analysis (omitted for now). Empirically, they show that the alignment metric for <code>SimCSE</code> is comparable to average BERT, but the uniformity measure is significantly lower, leading to much better performance in terms of accuracy. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="weng-2021---contrastive-representation-learning"><a class="header" href="#weng-2021---contrastive-representation-learning">Weng 2021 - Contrastive Representation Learning</a></h1> Reference: Weng, Lilian. (May 2021). Contrastive representation learning. Lil’Log. <a href="https://lilianweng.github.io/posts/2021-05-31-contrastive">https://lilianweng.github.io/posts/2021-05-31-contrastive</a>. This page just follows Lilian's notes on contrastive learning. Contrastive learning seeks to learn representations such that related samples are close together and unrelated samples are far apart. A sample is typically represented by a 1D vector (or embedding). The literature on contrastive learning comes from a wide range of tasks such as semantic retrieval, computer vision, hence we may see the same ideas repeated in slightly different forms. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="li-2021---taobao-embedding-based-retrieval"><a class="header" href="#li-2021---taobao-embedding-based-retrieval">Li 2021 - TaoBao Embedding-Based Retrieval</a></h1> <a href="https://arxiv.org/abs/2106.09297">Li 2021 - Embedding Based Product Retrieval in TaoBao Search</a> This paper explains TaoBao's embedding-based search retrieval system. Some interesting aspects of this paper: <ul> <li>Tackles the inherent trade-off between recall and relevance for embedding based retrieval</li> <li>Shows a way to incorporate personalization into search retrieval</li> <li>Ideas on enhancing representation of the search query</li> </ul> <h2 id="context"><a class="header" href="#context">Context</a></h2> TaoBao's product search has many stages. This paper focuses on the initial retrieval stage, when around <code>10k</code> products are retrieved per query. The retrievers used are: <ul> <li>Traditional lexical based term match retriever</li> <li>Item-based collaborative filtering retriever (presumably based on recent items)</li> <li>Embedding based retriever</li> </ul> The retrieval results from each retriever are merged and de-duplicated, and then sent to participate in the next phase (pre-ranking). This paper is only focused on the embedding-based retriever. <h2 id="problem"><a class="header" href="#problem">Problem</a></h2> Taobao tackles several problems in this paper: <ul> <li>Poor control of relevance for embedding based retrieval (EBR). Taobao reported that over time, the EBR contributed to an increase in complaints as non-relevant items started to get surfaced. For example, if a query is made for <code>Nike Shoes</code>, shoes from similar brands like <code>Adidas</code> may get surfaced.</li> <li>Balancing relevance and personalization. Taobao notes that personalization is an important aspect of product search, and they propose a way of merging search relevance and personalization in this paper.</li> </ul> <h2 id="setup-13"><a class="header" href="#setup-13">Setup</a></h2> Let U={u1,...,uN} denote a set of users, Q={q1,...,qN} denote their corresponding queries, and I={i1,...,iM} denote the collection of items. Let us divide the user u's historical sequence of items into 3 buckets: <ul> <li>Real time, i.e. most recent T item interactions, denote as Ru={i1u,...,iTu}</li> <li>Short term, i.e. within last <code>10 days</code> but older than real time, denote as Su={i1u,...,iTu}</li> <li>Long term, i.e. within last <code>1 month</code> but older than short term, denote as Lu={i1u,...,iTu}</li> </ul> The search task is as follows. We are given the historical behaviours (Ru,Su,Lu) of a given user and the submitted query qu at time t. We also get access to a sequence of historical queries for the user qhis={q1u,...,qku}. Our task is to return a set of items i∈I that satisfy the search request. Typically, we score each <code>query, item</code> pair according to score z and return the top-K items: z=F(ϕ(qu,Ru,Su,Lu), ψ(i)) Where: <ul> <li>F denotes the scoring function, typically the inner product function</li> <li>ϕ denotes the query + behaviour encoder</li> <li>ψ denotes the item encoder</li> </ul> <h2 id="model-architecture"><a class="header" href="#model-architecture">Model Architecture</a></h2> The model architecture in the paper is quite complicated, although understandable after some time. I don't think it is profitable to describe the architecture in full specificity, but there are useful ideas within. At a high level, the model architecture is a standard two-tower setup, where the user query and behaviours are encoded into a user embedding, and the item is encoded into an item embedding. The dot product is used to score user-item pairs and ANN search is used for retrieval. The interesting details lie mostly within the user tower, in particular the way the query and behaviours are integrated. <h3 id="query-representation"><a class="header" href="#query-representation">Query Representation</a></h3> We first tackle the Query Representation. Given a query string, the query encoder encodes it into Qmgs∈R6×d, where d is an arbitrary embedding dimension. The first dimension is <code>6</code> because the paper uses <code>6</code> different ways to encode the query (they call this multi-grained representation). The <code>6</code> representations are briefly explained as follows: <ul> <li>1-gram and 2-gram embeddings. The query string is tokenized into 1-gram or 2-gram and embeddings are looked up for each token. Mean pooling over tokens is done to get a single embedding.</li> <li>Phrase embeddings. Similarly, the query string is tokenized into phrases using their query segmentation engine. For example, 红色 is a phrase and 连衣裙 is a phrase. The same embedding lookup and mean-pooling is done. Call this qseg embedding.</li> <li>Phrase transformer embeddings. The phrase embeddings are also passed through a transformer before mean-pooling. I suppose the sequence of phrases matters enough to do this.</li> <li>Historical query embeddings. This is the interesting part where the query interacts with historical queries. <ul> <li>Specifically, the phrase embedding qseg∈Rd and the historical query matrix qhis∈Rk×d is used to form an attention matrix in the form of attn=softmax(qseg⋅qhisT)∈R1×k, which provides the relevance of each historical query to the current submitted query.</li> <li>The attention weights are then used to do a weighted average over the historical query embeddings qhis_seq=attn⋅qhis∈Rd.</li> <li>Hence we get a weighted representation of historical queries where more relevant queries to the current query are emphasized.</li> </ul> </li> <li>Mix embeddings. This is simply a sum over the above <code>5</code> embeddings</li> </ul> Note that in the query representation, both token-based embeddings and phrase-based embeddings are used to provide a fine grained representation of the query. <h3 id="user-behaviour-representation"><a class="header" href="#user-behaviour-representation">User Behaviour Representation</a></h3> Now we tackle the representation of user behaviours. There are some minor deviations on how they treat Ru,Su,Lu respectively, but broadly they follow the same structure. We refer to Ru below but similar treatment applies to Su or Lu. <ul> <li>Each item in the sequence Ru is represented by an embedding</li> <li>The item embedding comprises a concatenation of several embeddings, such as item ID embedding, item category embedding, brand embedding etc.</li> <li>Ru∈RT×d is passed into a transformer with self attention to get a hidden representation Rself_attnu∈RT×d</li> <li>Similar to the treatment of historical query above, cross attention is performed with the query representation qmgs to get a query-weighted representation of Ru, as follows: <ul> <li>The attention matrix is formed by taking attn=softmax(qmgs⋅(Rself_attnu)T)∈R6×T, which provides the relevance of each historical interacted item to the current submitted query.</li> <li>The attention weights are then used to do a weighted average over the historical item embeddings HR=attn⋅Rself_attnu∈R6×d.</li> <li>Hence we get a weighted representation of historical item interactions where more relevant items to the current query are emphasized.</li> </ul> </li> <li>The same operations are applied to Su and Lu to get HS and HL respectively.</li> </ul> Finally, the query qmgs and the representations of historical item interactions HR,HS,HL are passed into a self attention layer to get the final user representation, which captures both query semantics and historical information about the user: Hqu=Transformer([[CLS],qmgs,HR,HS,HL])∈R25×d The embedding at the [CLS] token is taken to represent the user. This has dimension d. <h3 id="item-representation"><a class="header" href="#item-representation">Item Representation</a></h3> The item embedding is represented similarly to how each item was represented above. Mean pooling is done over the phrase segmentation of the item title, and added to the ID embedding of the item. <h2 id="loss-function"><a class="header" href="#loss-function">Loss Function</a></h2> Now that we have described the model architecture, we proceed to training details. The authors used sampled softmax loss to train the model. The sampled softmax loss was compared against pairwise hinge loss and found to be superior. The hinge loss also has a downside of needing to tune the margin parameter, which can significantly affect results. Specifically, for a given user representation Hqu and item representation Hi of item i. Let the predicted score for the match between user and item i be denoted: y^(i∣q,u)=∑i′∈Iexp(F(Hqu,Hi′)/τ)exp(F(Hqu,Hi)/τ) Where τ is the temperature parameter. We can then denote the loss as the negative log likelihood of the positive interactions: L=−(q,u,i+)∈D∑logy^(i+∣q,u) <h2 id="adjustments"><a class="header" href="#adjustments">Adjustments</a></h2> The authors propose two adjustments to the training procedure to improve performance. Firstly, they propose to use temperature τ to handle noisy relevance labels. They argue that the positive signals in e-commerce data from clicks and purchases are noisy, and will benefit from tuning the temperature. The argument is: <ul> <li>As τ→0, the softmax function approaches a one-hot function of the item with the highest score. Thus the model will be encouraged to fit the training signals exactly, which runs the risk of overfitting to noisy labels.</li> <li>Conversely, as τ→∞, the softmax function approaches a uniform distribution over all items regardless of relevance. This encourages the model to underfit the training signals, since scoring positive items higher does not affect the loss by much.</li> </ul> Therefore, supposing the positive signals to be noisy, we should increase τ to accommodate the noise. The authors found that τ=3 performed the best for their data. <blockquote> Note: This approach advocates for more smoothing of the softmax distribution, which is opposite to the findings from papers like <a href="papers/./gao_2021.html">SimCSE</a>, which advocate for less smoothing (or a very low τ=0.05). I suppose the optimal temperature depends on how noisy the positive signals are relative to how noisy the negative signals are. </blockquote> Secondly, they propose a unique way of mining hard negative samples. Firstly, they mine hard negatives from the random negatives as per normal by choosing the top N negative items with the highest inner product with the user representation. Let this matrix be called Hhard∈RN×d. Now, they linearly interpolate Hhard with the positive item embedding Hi+ as follows: Hmix=α⋅Hi++(1−α)⋅Hhard This Hmix is then used in the loss function as the negative item embeddings. In their experiment, they found that α between 0.4−0.6 worked well. <blockquote> Note: This approach seems to exacerbate the problem of false negatives, since we are using the positive item as part of the negative signal. </blockquote> Finally, the authors propose a simple solution to the problem they raised of poor relevance matches surfaced by the EBR approach. The propose relevance control measure is to simply add a boolean term-based filter to the EBR results. For example, if the query is <code>Red Nike Shoes</code>, the boolean query will be something like <code>Color:Red AND Brand:Nike AND Product:Shoes</code>. The boolean filter will thus filter out any irrelevant results. The paper does not elaborate much on how the query filter is implemented. Taobao decides on the key query terms to filter based on their understanding of what constitutes essential components of the query (i.e. brand name, color, style etc.). <blockquote> Note: One may wonder what is the point of Embedding Based Retrieval (EBR) if we are going to slap a lexical boolean filter on it at the end. The answer is that Taobao is not using EBR as a way to achieve fuzzy semantic matching, as one might do when the lexical retrieval is too strict. Rather, it seems that due to the large number of items in Taobao's catalogue, the number of items passing the boolean filter already far exceeds the <code>10k</code> limit. Thus they are using EBR as a way of surfacing the most relevant items within the subset of items that pass the boolean filter. </blockquote> <h2 id="results-2"><a class="header" href="#results-2">Results</a></h2> Interestingly, Taobao evaluates the quality of the retrieval based on two metrics: <ul> <li>Recall@1k: This is a typical measure of how many positive items were retrieved by the algorithm</li> <li>Relevant %: Taobao has an in-house relevance model that scores <code>0.915</code> AUC on human-annotated data on whether a query and product has good relevance match. Hence they use this model to determine if each query and retrieved product has good relevance, and reports the % of good relevance.</li> </ul> The findings of the paper are: <ul> <li>The multi-grained representation of the query adds a small amount of recall</li> <li>Tuning temperature τ / using hard negatives adds significant amount of relevance</li> <li>Adding all the improvements raises relevance but decreases recall slightly, showing that there is some inherent trade-off between relevance and recall (i.e. some users do indeed engage with non-relevant items to their query)</li> </ul> It seems that Taobao prioritizes relevance over recall, which makes sense as surfacing irrelevant results undermines the integrity of the search system. Furthermore, since they use a boolean filter for relevance control, irrelevant results surfaced by the EBR system would get filtered out anyway. <h2 id="takeaways-3"><a class="header" href="#takeaways-3">Takeaways</a></h2> This paper tackles an important problem for EBR systems where irrelevant results get surfaced sometimes. While a boolean filter might make sense for Taobao's use case, other systems might need a different approach to avoid being overly strict in relevance control. It is also insightful to understand the inherent trade-off between recall and relevance when training the model. Finally, the paper also demonstrates a way to incorporate personalization into search via cross attention mechanisms between the query representation and the historical item interactions of the user. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="zou-2021---plm-based-ranking-in-baidu-search"><a class="header" href="#zou-2021---plm-based-ranking-in-baidu-search">Zou 2021 - PLM Based Ranking in Baidu Search</a></h1> <a href="https://arxiv.org/abs/2105.11108">Pre-trained Language Model based Ranking in Baidu Search</a> This paper explains how Baidu fine-tunes pre-trained language models (PLMs) to perform efficient ranking for search. Some interesting aspects: <ul> <li>A quick sentence extractor to perform query-sensitive summarization of documents on the fly</li> <li>Calibrating relevance signals for each query-document pair using click features</li> </ul> <h2 id="background-4"><a class="header" href="#background-4">Background</a></h2> The paper highlights some challenges with using PLM-based cross encoder ranking for search: <ul> <li>Documents tend to be long, making it hard for a ranker to meet the extremely low latency requirements of search. Transformer based models have quadratic time complexity with sequence length, exacerbating the problem.</li> <li>PLMs are typically trained on the language modelling task, which does not transfer directly into a query-document pair type of sentence for ranking.</li> <li>The score needs to be in a meaningful range to be easily blended with other ranking signals such as freshness, authority etc. Note: This portion might be specific to Baidu as they score documents in the <code>0-4</code> range.</li> </ul> <h2 id="problem-formulation"><a class="header" href="#problem-formulation">Problem Formulation</a></h2> Given a query q and a set of retrieved documents D={di}i=1N, we desire a scoring function f:Q×D→R, which maximizes some evaluation metric: maxf Eq,D,Y Θ(Y,F(q,D)) Where: <ul> <li>Θ is an evaluation metric like DCG</li> <li>F(q,D)={f(q,di)}i=1N is the set of document scores for this query</li> <li>Y={yi}i=1N is the set of true labels. For Baidu, yi is in the range <code>0-4</code>, corresponding to <code>bad</code>, <code>fair</code>, <code>good</code>, <code>excellent</code>, <code>perfect</code> respectively.</li> </ul> In learning to rank, we have a set of labelled query document pairs denoted as Φ={ϕq}, where ϕq={q,D={di},Y={yi},i=1,...,N} is a set of labelled documents given a query q. We minimize the learn to rank loss L as a proxy to minimizing the non-differentiable metric Θ: L(f)=∣Z∣1{q,D,Y}∈Φ∑l(Y,F(q,D)) Where l is the loss function and Z is the normalizing factor. <h2 id="query-weighted-summary-extraction"><a class="header" href="#query-weighted-summary-extraction">Query-Weighted Summary Extraction</a></h2> As mentioned above, documents are typically long and ill-suited for direct inputting into a BERT-based transformer. Thus the authors propose a fast query-sensitive algorithm to select a subset of sentences from each document to serve as the <code>summary</code>. The algorithm is as follows: <ul> <li>Given a query q, document d, decay factor α and max number of sentences k</li> <li>Tokenize q into a set of words Wq</li> <li>Tokenize the document d into a set of sentences S</li> <li>Initialize a vector of word importances as ωw for w∈Wq</li> <li>For each sentence si∈S: <ul> <li>Tokenize si into a set of words Wsi</li> <li>Compute score=∑w∈Wsi⋂Wq ωw</li> </ul> </li> <li>Choose the sentence with the highest score s∗ and add to summary</li> <li>Remove s∗ from S</li> <li>Decay used word importances in ω accordingly, i.e. ωw←α⋅ωw for all w∈Ws∗⋂Wq</li> </ul> It is a simple and fast algorithm that chooses sentences with the highest word overlap with the query. Each subsequent sentence aims to capture a different set of words by decaying the weights of words already chosen. Baidu chooses α=0.5 and eventually k=1 in their experiments. <h2 id="ranking-architecture"><a class="header" href="#ranking-architecture">Ranking Architecture</a></h2> The ranking architecture is a simple cross-encoder set up, except that the representation is split up with <code>query, title</code> on one side and the query-sensitive document summary on the other side. This avoids incurring quadratic time complexity on the fully concatenated text. In their experiments, they used <code>9</code> representation layers and <code>3</code> interaction layers, so that the full attention is only done for <code>3</code> layers. Baidu estimates that this setup reduces average time by <code>30%</code>. The representation at the <code>[CLS]</code> token is taken and probably attached with a softmax head to perform classification into the <code>0-4</code> range. <div class="table-wrapper"><table><thead><tr><th style="text-align: center"><img src="papers/../images/zou_2021_baidu_architecture.png" alt="Ranking Architecture" /></th></tr></thead><tbody> <tr><td style="text-align: center">Ranking Architecture</td></tr> </tbody></table> </div> The model is fine-tuned (the paper calls it pre-training, but I think it's more of a fine-tuning step) using triplet loss on a set of positive and negative pairs. The generation of positive and negative examples is covered below. Specifically, the PLM is fine-tuned to minimize the following loss: l(Y,F(q,D))=yi<yj∑[f(q,di)−f(q,dj)+m]+ Where yi,yj refer to the true labelled score of documents i,j respectively, and m is the margin. <h2 id="relevance-score-calibration"><a class="header" href="#relevance-score-calibration">Relevance Score Calibration</a></h2> Naively, we can curate yi,yj using clicks and non-clicks respectively. However, the paper highlights several issues with this approach: <ul> <li>Many clicks are in fact false positives, as they are caused by noisy clicks such as clickbait and accidental clicks</li> <li>There is exposure bias in the ranking system favouring higher position items</li> <li>Clicks do not necessarily imply relevance</li> </ul> Thankfully, there are many other signals accompanying each <code>q, d</code> pair, assuming that sufficient traffic has been observed for such a pair. Some examples from the paper include: <ul> <li># skip# click, which measures the number of clicks vs number of users who were impressed but did not click.</li> <li># total clicks in the query# click, which measures the number of clicks on the item as a fraction of all other clicks the user made for that query. This helps to control for users who are just trigger-happy and click on many items.</li> <li># click# long click, which measures the number of long clicks (i.e. user clicked and made no other action for <code>x</code> seconds) against clicks.</li> </ul> These features contain signals about relevance that can be combined in a relevance model to infer the relevance of a <code>q, d</code> pair. To do this, Baidu crowd-sourced annotated ratings of a score from <code>0-4</code> for <code>70k</code> query-document pairs. A simple decision tree-based model was then trained to infer the relevance score for a <code>q, d</code> pair given the relevance signal features discussed above. After training this model, it is then used to run inference on all <code>q, d</code> pairs (with sufficient observed traffic) to predict the relevance label and create a training dataset with a far more accurate relevance label. Note: the paper also has a final step of fine-tuning the model according to crowd-sourced data, after the initial fine-tuning phase above with the calibrated relevance scores. Since this part is not particularly innovative nor feasible, I have omitted it here. <h2 id="experiments-4"><a class="header" href="#experiments-4">Experiments</a></h2> They compared these innovations in an ablative manner from the original ERNIE system, which was a 12-layer transformer ranker trained using pairwise loss with human-labelled query-document pairs. The online ablation experiments show that: <ul> <li>The inclusion of the document summary in the document representation adds around <code>0.65%</code> in DCG@2</li> <li>The inclusion of fine-tuning on calibrated relevance scores increases the gains to <code>2.78%</code></li> </ul> The ablation experiments also showed that calibrated relevance scores is critical for effective fine-tuning. By using raw user clicks as signal, the positive-negative ratio (PNR, i.e. ratio of concordant <code>q,d</code> pairs over number of discordant <code>q,d</code> pairs) was only <code>1.86</code>, but with the calibrated scores it increased dramatically to <code>3.35</code>. <h2 id="takeaways-4"><a class="header" href="#takeaways-4">Takeaways</a></h2> This is an interesting paper on building an effective pure search ranker (i.e. no personalization), although we can always layer additional personalization on top of the initial relevance score. The benefit of training a pure search ranker is that we can observe multiple instances of the same <code>q, d</code> pair, which allows us to bootstrap a diverse array of weaker relevance signals into a much stronger relevance score. Even if we do not have human-annotated signals to train the decision tree, we often have a sparser but stronger signal (e.g. conversion) and many other weaker but more abundant signals. Hence, it may be possible to use the same approach to predict the stronger signal from the weaker ones. I suppose another way to deal with the diverse signal problem is multi-task learning, but this approach allows us to avoid the complexity of multi-task learning by first pre-processing relevance signals. The idea of a on-the-fly query-sensitive document summary also sounds useful, and can probably be implemented with the help of ElasticSearch. We could extend this approach by performing query expansion before doing the document summary, which may allow us to fetch more relevant sentences which do not have exact term overlap. Another way of doing this is to pre-generate document summaries using an LLM, which requires assuming that one summary for a document is able to suit all queries, which may not be the case. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="dao-2022---flash-attention"><a class="header" href="#dao-2022---flash-attention">Dao 2022 - Flash Attention</a></h1> <a href="https://arxiv.org/pdf/2205.14135.pdf">Paper Link</a> This paper argues that the attention mechanism is slow because of reading / writing between GPU High Bandwidth Memory and GPU on-chip SRAM. The authors hence create a block-wise attention algorithm that minimizes such IO read / writes and speeds up attention significantly especially when the sequence length is long. <h2 id="brief-overview-of-attention"><a class="header" href="#brief-overview-of-attention">Brief Overview of Attention</a></h2> Suppose we have an input sequence of embeddings X=(x1,...,xN) where xi∈Rk, such that X∈Rk×N. Naively, we can compute activations by V=XT⋅Wv, where Wv∈Rk×d, such that V∈RN×d. However, this naive way of encoding our input sequence does not allow interaction between inputs at different positions (say xi with xj). We can see this by observing that the first row of V is only affected by the first column of X (i.e. first encoding x1), and likewise for all the other positions. Attention addresses this problem by adding an interaction mechanism. Besides Wv, we also create weight parameters Wq,Wk∈Rk×d. Given an input X, we compute Q,K,V∈RN×d as follows: <ul> <li>Q=XT⋅Wq</li> <li>K=XT⋅Wk</li> <li>V=XT⋅Wv</li> </ul> We then create an interaction matrix S=QKT∈RN×N, and apply row-wise softmax to get P=softmax(S)∈RN×N. S can be thought of as a pairwise similarity matrix between the encoding at position i and position j that captures the degree of interaction. For example, in a sentence <code>the economy has been in decline</code>, the value of S1,5 (assuming 0-index) measuring the interaction between <code>economy</code> and <code>decline</code> might be high. Finally, we produce the output O=PV∈RN×d, which is an activation output from the input sequence that has captured the interactions between tokens at different positions of the input. This simple mechanism has led to significant improvements in language modelling. <h2 id="gpu-memory-hierarchy"><a class="header" href="#gpu-memory-hierarchy">GPU Memory Hierarchy</a></h2> The memory hierarchy is such that read/write speed is super fast on the SRAM but memory is highly limited. Hence, the N x N attention matrix is written/read repeatedly to/from HBM, resulting in IO being a bottleneck. The numbers are as such on an A100 GPU: <ul> <li>SRAM: 19 TB/s (20 MB RAM)</li> <li>HBM: 1.5 TB/s (40 GB RAM)</li> </ul> <h2 id="naive-attention-algorithm"><a class="header" href="#naive-attention-algorithm">Naive Attention Algorithm</a></h2> The naive attention algorithm has many reads and writes to HBM. (ps: Not sure why we cannot persist the intermediate matrices on SRAM and complete the computations, but in any case the naive algorithm requires materializing the N×N matrices on SRAM which will quickly flood it. For example, a sequence length of <code>2,048</code> at <code>float32</code> already takes up <code>33MB</code> for the S matrix). <ol> <li>Load Q,K from HBM, compute S=QKT, write S to HBM</li> <li>Read S from HBM, compute P=softmax(S), write P to HBM</li> <li>Load P,V from HBM, compute O=PV, write O to HBM</li> </ol> <h2 id="flash-attention-1"><a class="header" href="#flash-attention-1">Flash Attention</a></h2> The main idea is quite simple: instead of computing the full attention matrix, we use block-wise tiling to compute parts of it at a time. This reduces the memory required for each block and allows the whole computation to be done on SRAM while minimizing the amount of IO read from HBM, leading to faster compute time and lower memory usage on SRAM. The difficulty is in devising a block-wise softmax algorithm that yields the exact same result as computing it all at once. Consider the naive softmax algorithm on an arbitrary vector x∈RB. m(x)p(x)l(x)softmax(x):=imax xi:=[exi−m(x)...exB−m(x)]:=i∑p(x)i:=l(x)p(x) Note that the maximum value m(x) is subtracted for numerical stability to avoid overflow (underflow is ok because e−inf=0). f(x)∈RB is the numerator and l(x)∈R is the sum of all elements in p(x). Now, the problem with the naive softmax algorithm in the context of attention is that we need an entire row of S (N elements) to perform the row-wise softmax computation. This will not be available if we are performing block-wise computation, since we are splitting Q,K∈RNxd row-wise into blocks of Qi,Kj∈RBxd. When we compute Sij:=Qi⋅KjT∈RBxB, blocks of S will be materialized in each pass, but not the entire row at a time. Hence, we need a modified algorithm that allows us to compute chunks of the final output O∈RNxd at a time by iterating block-wise through S, such that the combination of the new chunk of O at each step with the already written intermediate O gives the correct result at the end. The key to realizing this algorithm is in decomposing the softmax step, as shown below. Consider two vectors xa,xb∈RB. We can decompose the softmax of their concatenated vector x=[xa xb]∈R2B as follows: m(x)p(x)l(x)softmax(x)=max(m(xa),m(xb))=[em(xa)−m(x)⋅p(xa)em(xb)−m(x)⋅p(xb)]=em(xa)−m(x)⋅l(xa)+em(xb)−m(x)⋅l(xb)=l(x)p(x) The first line of the above simply notes that the maximum of x is the maximum over each of the subvector maximums xa,xb. The second line notes that we previously multiplied each element of p(x) by a factor, say e−m(xa)⋅exi for those in xa. To get the correct multiplier for the full vector x, we need to divide away the previous multiplier and apply the new multiplier, i.e. e−m(x)/e−m(xa)=em(xa)−m(x). The third line notes that the new denominator is the sum over each of the subvector sums, after we apply the correct multiplier from line 2. The decomposition is simple but powerful. It implies that so long as we keep track of intermediate statistics m(x) and l(x), we can compute the softmax of a long vector x by splitting x into subvectors and operate over each subvector at a time. Now we are ready for Algorithm 1: Flash Attention of the paper. <hr /> Require: Q,K,V∈RNxd01. Initialize O=0Nxd,l=0N,m=−infN02. Divide Q,K,V row-wise into B-sized blocks of Qi,Kj,Vj∈RBxd03. Divide O row-wise into B-sized blocks of Oi∈RBxd04. Divide m,l into B-sized arrays of mi,li∈RB05. for j in 1,... :06. Load Kj,Vj from HBM to SRAM07. for i in 1,... :08. Load Qi,Oi,li,mi from HBM to SRAM09. Compute Sij=Qi⋅KjT∈RBxB10. Compute mij=rowmax(Sij)∈RB11. Compute Pij=eSij−mij∈RBxB12. Compute lij=rowsum(Sij)∈RB13. Compute minew=elementwise-max(mi,mij)∈RB14. Compute linew=emi−minew⋅li+emij−minew⋅lij∈RB15. Write Oi←(lnew)−1×[li×emi−minew×Oi+emij−minew×PijVj] to HBM16. Write li←linew,mi←minew to HBM17. Return O(1)(2)(3)(5)(7) <hr /> Note that we use 0Nxd,0N to denote a zero matrix of size Nxd and a zero array of size N respectively. For simplicity, we divide Q,K,V into equal B-sized blocks but the paper allows different block sizes for Q and K,V. The equation numbers on the right in parentheses show which equations the lines correspond to above. Equation line 15 is a bit confusing because it combines multiple steps together. The next few paras try to unpack this. Firstly, note that we are using the × operator to denote an element-wise broadcasted multiplication. For a vector li∈RB and a matrices Pij∈RBxB,Vj∈RBxd, observe the associative property (li×Pij)⋅Vj=li×(Pij⋅Vj), since each element of li only affects the corresponding row in the final matrix. This allows us to apply the scaling to either Oi or Pij and the result will be the same. Next, see that the term emij−minew×PijVj is simply the corrected numerator of the softmax dotted with Vj. Dividing this term by lnew gives the output block for this particular Sij,Vj pair. Similarly, the other term li×emi−minew×Oi is the existing output that has been accumulated from previous steps {Sij,Vj:j=1,...}. Due to the associative property, we can also directly apply the scaling correction to Oi. The li/lnew×emi−minew are scaling factors according to equations (6), (8) to correct the scaling of previous steps. Finally, we should understand why there is a <code>+</code> in equation 15. I find it easier to visualize if we set B=1. If we trace the matrix multiplications, we will observe that Oi∈R1xd is only affected by Qi∈R1xd, i.e. it corresponds to only the query token in position i. Now, Oi represents the weighted average over all N positions of the V matrix where the weights are determined by the softmax over the interaction between Qi (representing one token) and all N positions on the K matrix. This weighted average is why it is a + symbol: we are accumulating the weighted sum over Vj into Oi. The only complication is that we are applying the scaling corrections at each step. Hopefully these explanations provide some intuition to the FlashAttention algorithm, which is quite a simple idea but makes a ton of difference practically. It should be easy to implement this algorithm in <code>numpy</code> if the reader wishes to understand it better. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="wei-2022---cot-prompting-in-llms"><a class="header" href="#wei-2022---cot-prompting-in-llms">Wei 2022 - CoT Prompting in LLMs</a></h1> <a href="https://arxiv.org/abs/2201.11903">Chain-of-Thought Prompting Elicits Reasoning in Large Language Models</a> This is an early paper that sparked the explosion of research on in-context learning for LLMs. The main idea is to use few-shot examples with the Chain-Of-Thought reasoning, i.e. <code><Question>, <CoT reasoning>, <Answer> </code>, as opposed to just the question and answer alone <code><Question>, <Answer></code> (as per <a href="https://arxiv.org/abs/2005.14165">Brown 2020 - Language Models are Few-Shot Learners</a>). <h2 id="method-3"><a class="header" href="#method-3">Method</a></h2> The method is simple - we include the CoT reasoning in the few-shot examples included in the prompt, inducing the LLM to also generate CoT reasoning before answering the question at inference time. An example prompt: <pre><code>Q: Roger has 3 balls. He buys 2 more cans of balls, each with 3 balls. How many balls does he have now? A: Roger started with 3 balls. 2 cans of 3 balls each is 6 balls. 3 + 6 = 9. The answer is 9. </code></pre> What are some advantages of CoT prompting? <ul> <li>Useable on black box LLMs. No fine-tuning is required, so we can readily apply it on off-the-shelf LLMs.</li> <li>CoT allows decomposition of complex problems into intermediate steps. The conceptual idea is that this allows the model to offload additional computation to the intermediate tokens, analogous to human reasoning.</li> <li>CoT reasoning offers interpretation of model behaviour. We have a suggestive trace of how the model arrived at that answer (though it would be naive to assume that the LLM uses the CoT trace reasoning exactly the way human logic operates).</li> </ul> <h2 id="observations"><a class="header" href="#observations">Observations</a></h2> An important sub-finding is that between CoT prompting and standard prompting (as per Brown 2020), CoT prompting's improvement gap significantly increases as we: <ul> <li>Increase the question difficulty (seems intuitive); and</li> <li>Increase the model size. At the <code>8B</code> parameter model size, the study showed no difference between CoT prompting and standard prompting. But the gap widened significantly at the <code>100B</code> model size and widened further at the <code>500B</code> model size.</li> </ul> Another interesting ablation study tried to isolate the improvements to a specific aspect of prompting: <ul> <li>Equation only. This would be something like <code>3 + 2 * 3 = 9</code> for the above example. This showed to be no different from standard prompting. This showed to be useful on some datasets with simpler steps, but did not help on GSM8K which requires more semantic parsing of the question.</li> <li>Variable compute only. One may argue that the exact tokens generated in the intermediate step does not matter much, all that matters is the additional compute the model performs to generate the intermediate tokens. Hence the authors prompt the model to generate dots <code>...</code> as the reasoning step instead. This proves to not be helpful.</li> <li>CoT after answer. Another argument is that including the CoT traces in the prompt improves in-context learning in and of themselves, meaning that the intermediate tokens are not actually necessary for improving the model's accuracy. The authors disprove this hypothesis by putting the CoT reasoning after the answer in the prompt, i.e. <code><Question, Answer, CoT reasoning></code>. This forces the LLM to generate the answer before it generates the CoT trace. This also proves to not be helpful.</li> </ul> Thus the ablation studies help to clarify that it is the intermediate natural language reasoning steps that help the model offload computation and improve the accuracy of its answers. <h2 id="results-3"><a class="header" href="#results-3">Results</a></h2> The performance of CoT few-shot prompting compared to standard few-shot prompting is striking: <ul> <li>Using <code>GPT-3 175B</code>, performance increases from <code>15.6</code> to <code>46.9</code></li> <li>Using <code>PaLM 540B</code>, performance increases from <code>17.9</code> to <code>56.9</code></li> </ul> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="honovich-2022---instruction-induction"><a class="header" href="#honovich-2022---instruction-induction">Honovich 2022 - Instruction Induction</a></h1> <a href="https://arxiv.org/abs/2205.10782">Instruction Induction: From Few Examples to Natural Language Task Descriptions</a> This is an early paper that suggests taking a random subset of training <code>(input, label)</code> pairs, showing them to an LLM, and asking the LLM to guess the instruction that produces the label from the inputs. <h2 id="method-4"><a class="header" href="#method-4">Method</a></h2> For example, the LLM prompt may look something like: <pre><code class="language-bash">Here are the input-output pairs: Input: As soon as you can. Output: At your earliest convenience. ... (+4 more input / output pairs) The instruction was ... </code></pre> The LLM may then guess something like <code>The instruction was translate the inputs into more formal language</code>. We can then use this instruction as the new optimized prompt. Note that this method is a one step process, i.e. it does not iterate for further improvements. But I suppose we can sample multiple random subsets of training instances and generate prompts for each before picking the best prompt based on execution accuracy on the validation set. <h2 id="evaluation-2"><a class="header" href="#evaluation-2">Evaluation</a></h2> The paper found that the execution accuracy of these prompts with instruction fine-tuned version of <code>GPT-3</code> (called <code>InstructGPT</code>) could reach similar levels to human-generated prompts on most simple tasks. Notably, the authors were careful to control for the variation caused by the selection of few-shot examples: <ul> <li>An <code>induce set</code> of examples are held out from the training set</li> <li>For each run, <code>5</code> input / output pairs are randomly drawn from the induce set, and used to generate a prompt</li> <li>The evaluation accuracy of this prompt is recorded</li> <li>The evaluation accuracy is averaged over <code>100</code> runs</li> </ul> In hindsight, by using the idea of <a href="papers/./wang_2022.html">self-consistency</a>, the performance can likely be significantly improved by taking the majority vote over the <code>100</code> runs rather than the average accuracy. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="huang-2022---llms-can-self-improve"><a class="header" href="#huang-2022---llms-can-self-improve">Huang 2022 - LLMs can Self Improve</a></h1> <a href="https://arxiv.org/abs/2210.11610">Large Language Models Can Self-Improve</a> The main idea of this paper is that we can improve the instruction tuning of LLMs for reasoning capabilities using its own synthetic generated data. <h2 id="method-5"><a class="header" href="#method-5">Method</a></h2> We are given a pre-trained LLM and a question-only training dataset (e.g. like GSM8k). We are also given a few-shot Chain of Thought examples (an example comprises of a question, a reasoning, and a correct answer). The method is simple. For each question xi in the training set, we: <ul> <li>Sample m reasoning paths and answers</li> <li>Use majority voting from the m answers to select the most consistent answer y~ (this is called self-consistency in the literature). Importantly, to increase diversity: <ul> <li>Set the temperature T>0</li> <li>Apply mixed formats of prompts and answers</li> </ul> </li> <li>Keep all reasoning paths that lead to y~ as our synthetic dataset</li> <li>Fine-tune our LLM on the synthetic dataset using supervised fine-tuning</li> </ul> Note that since we are using self-consistency to obtain "labels" for our synthetic dataset, we do not require training labels. <h2 id="observations-1"><a class="header" href="#observations-1">Observations</a></h2> For this method to work, self-consistency needs to be a reliable way of getting accurate answers. The authors plot the confidence score (% of paths leading to y<del>) of y</del> against the accuracy at that confidence level, and find that it is highly correlated. This implies that highly consistent answers is a strong indication of correctness. Generally performance increases as we increase the number of sampled paths m. It seems to saturate at around m=32. Also, the ideal temperature is around <code>1.2</code>, showing that diversity is important for this technique to work well. <h3 id="findings"><a class="header" href="#findings">Findings</a></h3> The fine-tuned LLM significantly advances the SOTA performance: <ul> <li>GSM8K increases from <code>74.4</code> using self-consistency to <code>82.1</code> (using self-consistency on the fine-tuned LLM)</li> </ul> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="tunstall-2022---setfit"><a class="header" href="#tunstall-2022---setfit">Tunstall 2022 - SetFit</a></h1> <a href="https://ar5iv.labs.arxiv.org/html/2209.11055">Paper link</a>. The task which SetFit tackles is few-shot labelling of texts. Since labelling data is scarce, current SOTA methods rely on large language models and techniques like In Context Learning (ICL) or Parameter Efficient Fine Tuning (PEFT). <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="wang-2022---self-consistency-llm"><a class="header" href="#wang-2022---self-consistency-llm">Wang 2022 - Self Consistency LLM</a></h1> <a href="https://arxiv.org/abs/2203.11171">Self-Consistency Improves Chain of Thought Reasoning in Language Models</a> The main idea of this paper is that majority decoding (dubbed self-consistency) of sampled LLM responses significant improves over CoT prompting (<a href="papers/./wei_2022.html">see Wei 2022</a>). <h2 id="method-6"><a class="header" href="#method-6">Method</a></h2> The idea is simple. As per (<a href="papers/./wei_2022.html">Wei 2022</a>), we use CoT few shot examples in our prompt. Instead of just picking the answer from one run of the LLM, we sample many <code><CoT reasoning, Answer></code> pairs from the LLM with the same prompt. We then "marginalize out" the reasoning by just choosing the answer that occurs the most frequently. Note that this method is unsupervised and the only cost to pay is the compute cost of multiple runs. The name self-consistency comes from the idea that the most consistent answer given by the model is the most reliable one. Prior approaches to task-based responses is to use greedy decoding, i.e. set T=0 and get the highest likelihood answer. This paper shows that an ensemble of diverse answers is far more effective than the greedy approach. We may think of this idea as analogous to random forests. Creating an ensemble over diverse weak learners improves significantly, and increasing the diversity of the learners (via column or row sampling) up to a certain point helps to improve performance. Note that the idea of self-consistency decoding is orthogonal to the specific choice of CoT prompting. Specifically, self-consistency significantly improves with several forms of CoT prompting, including: <ul> <li>Few shot in-context CoT prompts (<a href="https://arxiv.org/abs/2201.11903">Wei 2022</a>)</li> <li>Zero shot <code>Let's think step by step</code> prompt (<a href="https://arxiv.org/abs/2205.11916">Kojima 2022</a>)</li> </ul> <h2 id="parameters"><a class="header" href="#parameters">Parameters</a></h2> To generate diverse reasoning paths, the authors mainly refer to methods used in other papers: <ul> <li>For <code>PaLM-540B</code> they used <code>T=0.7</code> and <code>k=40</code> with top-k token truncation</li> <li>For <code>GPT-3</code> they used <code>T=0.7</code> without top-k token truncation</li> </ul> For the paper, the authors mainly sampled <code>40</code> paths. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="lee-2022---rq-vae"><a class="header" href="#lee-2022---rq-vae">Lee 2022 - RQ-VAE</a></h1> <a href="https://arxiv.org/abs/2203.01941">Autoregressive Image Generation using Residual Quantization</a> Vector quantization <a href="papers/./van_den_oord_2017.html">VQ-VAE</a> is used to represent an image as a sequence of discrete codes. After quantizing, an autoregressive model (AR model, e.g. transformer) is used to predict the codes. This paper aims to reduce the sequence length of the discrete codes, as it becomes more computationally efficient for the AR model. However, reducing the sequence length (aka reducing the spatial resolution) causes a rate-distortion trade-off. The dilemma is such: <ul> <li>Increasing the codebook size improves resolution of discretization to preserve quality</li> <li>But increasing the codebook size also increases the probability of codebook collapse, where only a subset of codes are used</li> </ul> The main idea of RQ-VAE is to reduce the spatial resolution by more precisely approximating the feature map at each location. Instead of increasing the codebook size, we use the codebook to recursively quantize the feature map in a coarse-to-fine manner. That is, the feature representation at each location is the sum of the selected codebook vectors at each level. For a codebook where each level has K codebook vectors and D levels, we can represent KD vectors for learning just D×K vectors. <ul> <li>The paper claims that due to this precise approximation, they can use just <code>8 x 8</code> latents to represent a <code>256 x 256</code> image. That is a reduction of <code>1024x</code>!</li> </ul> <blockquote> Note: other papers usually use a distinct set of vectors for each codebook level, but in this paper they use a shared codebook for all levels. This is probably an implementation detail / hyperparameter to be tuned. </blockquote> <h2 id="method-7"><a class="header" href="#method-7">Method</a></h2> There are two main stages: <ul> <li>Stage 1: Residual Quantized VAE.</li> <li>Stage 2: RQ-Transformer.</li> </ul> <h3 id="stage-1-residual-quantized-vae"><a class="header" href="#stage-1-residual-quantized-vae">Stage 1: Residual Quantized VAE</a></h3> Let a codebook C be a finite set {(k,e(k))},k∈[K] be a codebook of tuples of a code <code>k</code> and its code embedding e(k)∈Rnz. <ul> <li>K is the codebook size</li> <li>nz is the code embedding dimension</li> </ul> Given a vector z∈Rnz, let Q(z;C) denote the discrete code of z: Q(z;C)=k∈[K]argmin∣∣z−e(k)∣∣22 The normal VQ-VAE flow is such: <ul> <li>Start with an input image X∈RHo×Wo×3, where Ho,Wo are the original image dimensions</li> <li>The encoder extracts a feature map Z=Enc(X)∈RH×W×nz, where H=Ho/f,W=Wo/f are downsampled f times to form the latent feature map</li> <li>Applying the discretization independently to each position, we obtain a code map M∈[K]H×W: Mh,w=Q(Zh,w;C)</li> <li>We also obtain the quantized feature map Z^∈RH×W×nz: Z^h,w=e(Mh,w)</li> <li>Finally, we decode the quantized feature map to reconstruct the image: X^=G(Z^)</li> </ul> For Residual Quantization, we instead represent z∈RD as an ordered tuple of D codes: RQ(z;C,D)=(k1,...,kD)∈[K]D Where kd is the discrete latent code of z at depth d. Specifically, we obtain kd by recursively quantizing the residual. Starting with residual ro=z, we compute the next kd,rd as follows: kdrd=Q(rd−1;C)=rd−1−e(kd) Finally, we represent z^:=z^(D)=∑d=1De(kd) as the quantized embedding of z using the sum of embedding across all depths. <blockquote> Note: by using a shared codebook and summing up the embeddings at each depth, this resembles bloom hashing (see <a href="papers/./weinberger_2009.html">Weinberger 2009</a>). The difference is that we recursively / sequentially encode the residuals, whereas bloom hashing simultaneously applies multiple independent hashes and sums up those embeddings at the hashed positions. However, as the experiments indicate, the embedding norms get smaller as we go deeper into the codebook levels, since it learns a coarse-to-fine encoding. So it seems to make better sense to have separate codebook vectors for each level. </blockquote> The RQ-VAE flow is thus identical to the VQ-VAE flow above, except that the discretization step differs: <ul> <li>Start with an input image X∈RHo×Wo×3</li> <li>Encoder extracts Z=Enc(X)∈RH×W×nz</li> <li>Discretize each position to get code map M∈[K]H×W×D (note the additional dimension D, since we now have D codes per position): Mh,w,dZ^h,w=RQ(Zh,w;C,d)=d=1∑De(Mh,w,d)</li> <li>Finally decode: X^=G(Z^)</li> </ul> <h3 id="rq-vae-training"><a class="header" href="#rq-vae-training">RQ-VAE Training</a></h3> The loss is: LreconLcommitL=∣∣X−X^∣∣22=d=1∑D∣∣Z−sg[Z^(d)]∣∣22=Lrecon+βLcommit The first loss is the reconstruction error of making the RQ-VAE roundtrip. The second loss is the commitment loss, which aims to push the encoder to reduce the quantization errors. <ul> <li>Note that the commitment loss sums up the quantization error at each level of the codebook, not just the final quantization error at the end</li> <li>Also note that the authors did not include the loss to update the codebook as in <a href="papers/./van_den_oord_2017.html#learning">VQ-VAE</a>; instead they use the Exponential Moving Average approach to update the codebook C using the average encoder output assigned to each cluster</li> </ul> <h3 id="stage-2-rq-transformer"><a class="header" href="#stage-2-rq-transformer">Stage 2: RQ-Transformer</a></h3> Given that we have learned a code map of M∈[K]H×W×D, the goal is to now autoregressively predict the code map (laid out in raster scan order). Specifically, let us lay out M into a 2D array of codes S∈[K]T×D, where T=H×W. Thus each row of S (call it St) will contain D codes: St=(St1,...,StD)∈[K]D Note that S is laid out in raster scan order, meaning that we read the first row of pixels left-to-right, then move to the next row and so on. Thus the goal of stage 2 is to learn an autoregressive model which learns the joint probability function p(S), given all the prior tokens (i.e. causal masking). Note that we predict D codes for each position before moving on to the next: p(S)=t=1∏Td=1∏Dp(Std∣S<t,St,<d) The natural approach is to just fit an autoregressive transformer decoder to this unrolled sequence (of length T×D, where T=H×W). This incurs attention bottleneck complexity of T2D2. However, this paper argues that this is inefficient as it does not exploit the structure of our tokens, and proposes a customized RQ-transformer for this situation. The idea is that there are two orthogonal dimensions, so we can decouple and create a specialized transformer to handle each one orthogonally: <ul> <li>The first two dimensions, H and W, encode the position</li> <li>The last dimension, D, encodes the depth</li> </ul> <div class="table-wrapper"><table><thead><tr><th style="text-align: center"></th></tr></thead><tbody> <tr><td style="text-align: center"><img src="papers/../images/lee_2022_rqvae_rq_transformer.jpg" alt="RQ-Transformer Architecture" /></td></tr> <tr><td style="text-align: center">RQ-Transformer Architecture</td></tr> </tbody></table> </div> <h4 id="spatial-transformer"><a class="header" href="#spatial-transformer">Spatial Transformer</a></h4> The spatial transformer encodes the position dimension and marginalizes away the depth dimension. It is concerned with the "big picture". Specifically, each input position to the spatial transformer is: ut=PET(t)+d=1∑De(St−1,d) for t>1 Note that: <ul> <li>PET(t) is the position encoding for spatial position t</li> <li>We re-use the quantized embedding codebook vectors to represent inputs to the spatial transformer</li> <li>The depth position is marginalized out, such that we only need to encode T=H×W positions with the spatial transformer</li> <li>u1 is a special case for start of sequence and will have its own learnable embedding</li> </ul> The sequence of inputs (ut)t=1T is passed into a causally masked transformer to produce spatial context vectors ht like so: ht=SpatialTransformer(u1,...,ut) <h4 id="depth-transformer"><a class="header" href="#depth-transformer">Depth Transformer</a></h4> Given the context vector ht provided by the spatial transformer, the depth transformer predicts a very short sequence of D codes for this position t. Thus, the depth transformer can be a smaller stack of transformer layers. <ul> <li>Note: in the experiments, generally Nspatial=24 and Ndepth=4 in number of layers</li> </ul> Specifically, at spatial position t and depth d, input vt,d to the depth transformer is the sum of codebook embeddings up to depth d: vt,dvt,d=PED(d)+ht=PED(d)+d′=1∑d−1e(St,d′) for d=1 for d>1 The short sequence of inputs (vt1,...,vtD) is passed into a causally masked transformer + classifier head to predict the codes for each position. The same depth transformer is reused across all spatial positions t. The autoregressive loss is simply the negative log likelihood of the correct latent token sequence: LAR=ESEt,d[−logp(St,d∣S<t,St,<d)] <h4 id="computation-savings"><a class="header" href="#computation-savings">Computation Savings</a></h4> As observed above, fitting a naive transformer to the full sequence of length TD will incur attention complexity of T2D2. Compare this to the RQ-transformer: <ul> <li>Spatial transformer: T2, since we marginalize away the depth dimension</li> <li>Depth transformer: T×D2, since we re-use the depth transformer T times</li> <li>RQ transformer: T2+TD2</li> </ul> With T=512,D=4, we have around <code>15x</code> lower computational complexity: <ul> <li>Naive approach: <code>4,194,304</code></li> <li>RQ transformer: <code>270,336</code></li> <li>Ratio: <code>~15x</code>!</li> </ul> <h3 id="other-tricks"><a class="header" href="#other-tricks">Other Tricks</a></h3> There are two other tricks that they use in this paper, presumably to improve performance. The aim is to reduce the impact of exposure bias, which is the divergence between the samples during training and inference. During training, due to teacher forcing, we still get correct tokens to predict from at each position. During inference, since we generate each position autoregressively, errors can compound and lead the model to predict from a sequence that deviates far from what it has seen during training. The main idea is to introduce some uncertainty into both the labels and inputs to the training process. Let us define a categorical distribution on [K] conditioned by the encoder embedding z∈Rnz, with probability distribution: Qτ(k∣z)∝e−∣∣z−e(k)∣∣22/τ for k∈[K] Thus the probability of a discrete token k increases exponentially as its quantized embedding e(k) gets closer to the encoder embedding z. Note that as τ approaches 0, Qτ gets sharper and converges to the one-hot distribution Q0(k∣z)=1[k=Q(z;C)]. This soft approximation of the <code>argmin</code> is used for soft labelling of the objective when training the RQ-transformer. Specifically, note that at position t and depth d, we quantize by doing a nearest neighbour search for the residual vector rt,d−1. Hence the hard target for the RQ-transformer training is the one-hot label Q0(⋅∣rt,d−1). Instead we replace with Qτ(⋅∣rt,d−1). We also do Stochastic Sampling of the latent codes when generating training samples for the RQ-transformer. Instead of deterministic code selection when encoding the raw images to latent codes, we sample from the categorical distribution of Qτ(⋅∣rt,d−1). The ablation studies show that using soft labelling + stochastic sampling significantly improves performance. <h2 id="results-and-hyperparameters"><a class="header" href="#results-and-hyperparameters">Results and Hyperparameters</a></h2> General hyperparameters: <ul> <li>Codebook size: K=16,384</li> <li>Codebook levels: D=4</li> <li>Latent dimension size: nz=256</li> <li>Temperature: τ=0.5</li> </ul> Generally RQ-VAE matches or outperforms image generation of other autoregressive competitors like VQ-GAN. It is also a lot faster in generation speed and memory usage. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="tay-2022---differentiable-search-index"><a class="header" href="#tay-2022---differentiable-search-index">Tay 2022 - Differentiable Search Index</a></h1> <a href="https://arxiv.org/abs/2202.06991">Transformer Memory as a Differentiable Search Index</a> This paper proposes a new paradigm of search. Instead of using Approximate Nearest Neighbour (ANN) search using embeddings, the paper proposes to train a transformer seq2seq model to predict the ID of a document given a search context. Although the paper is focused on search and Question-Answering, the approach is general and can be used for any recommendation setting. <h2 id="comparison-to-traditional-search-index"><a class="header" href="#comparison-to-traditional-search-index">Comparison to Traditional Search Index</a></h2> Traditional search index such as BM25 or ANN search has three steps: <ol> <li>Indexing. Indexing involves storing representations of each document to enable easy search later on. For traditional lexical search like BM25, a sparse and high dimensional vector is stored for each document and an inverted index is built. For ANN search, a dense vector is stored and an index like HNSW is built.</li> <li>Training. For lexical search, the only "training" is to store the TF-IDF quantities in the index. For ANN search, a contrastive learning objective is used to shape the embedding space for effective retrieval.</li> <li>Retrieval. For lexical search, a lookup of the inverted index based on the query terms is performed. For ANN search, a Maximum Inner Product (MIPS) search is performed using some algorithm.</li> </ol> The differentiable search index operates differently. The authors argue that it is more elegant: <ol> <li>Indexing. Indexing is just a specific form of model fine-tuning. Specifically, the task is to predict the <code>docid</code> given the <code>content</code> of the document, i.e. <code>f(doc_content) = docid</code>. <ul> <li>There is no need to store a separate search index as the only artifact we need becomes the transformer model itself</li> </ul> </li> <li>Training. Training is a form of supervised fine-tuning where we predict <code>f(query) = relevant docid</code>.</li> <li>Retrieval. Retrieval is simply running a forward pass on the transformer on a new query.</li> </ol> The nice thing about this setup is that it extends naturally to a multi-task setting, by using different prefixes for different tasks. This is similar to the T5 approach for multi-task learning. <h2 id="ablations"><a class="header" href="#ablations">Ablations</a></h2> The main idea is simple and intuitive, but there are many details to the implementation that are non-trivial. <h3 id="indexing-strategy"><a class="header" href="#indexing-strategy">Indexing Strategy</a></h3> The indexing step aims to teach the transformer model how to link each <code>docid</code> with the contents of the document. This is entirely new knowledge to the transformer since the docids are arbitrary constructs. The following strategies were explored: <ol> <li><code>f(document context) = docid</code>: this had the best performance</li> <li><code>f(docid) = document context</code>: this performed terribly</li> <li><code>Mix of 1. and 2.</code>: Mix training batches with both options <code>1</code> and <code>2</code>, with a prefix string to indicate which direction.</li> </ol> <h3 id="indexing-and-training-order"><a class="header" href="#indexing-and-training-order">Indexing and Training Order</a></h3> Recall that we have two modes: <ul> <li>In indexing mode, our training data is <code>f(document content) = docid</code></li> <li>In training mode (or retrieval task), our training data is <code>f(query) = relevant docid</code></li> </ul> The following strategies were explored: <ol> <li><code>Index first, train second</code>. This means that we train with the indexing mode until convergence, then switch to training mode</li> <li><code>Multi-task</code>. This means that we mix training batches of both indexing mode and training mode and just use a prefix to indicate the different task setting</li> </ol> The <code>multi-task</code> strategy worked much better. The performance was sensitive to the mix between indexing and retrieval task. The ideal mix was found to be around <code>32</code> instances of indexing task to <code>1</code> instance of retrieval task. <h3 id="representation-of-document-ids"><a class="header" href="#representation-of-document-ids">Representation of Document IDs</a></h3> The following strategies were explored: <ol> <li> <code>Arbitrary Atomic Identifiers</code>. This means assigning an arbitrary random identifier for each document, and letting each identifier be a new token in the vocabulary. The downside is that there is an explosion in the vocabulary if we have a lot of items </li> <li> <code>Semantically Structured IDs</code>. This means representing each document as an ordered list of tokens, where the tokens contain some semantic meaning. This is implemented as a hierarchical k-means clustering. <ul> <li>Firstly, each document is embedded using a universal T5 embedder.</li> <li>Next, k-means with <code>10</code> centroids is performed. The cluster ID <code>0-9</code> is assigned as the first document ID token for each document</li> <li>Each cluster is then further clustered to generate the subsequent document ID tokens</li> </ul> </li> </ol> <h2 id="inference-time"><a class="header" href="#inference-time">Inference Time</a></h2> To generate recommendations at inference time, given a query, we pass it into the transformer model and just do generative decoding with beam search. Beam search maintains at each decoding step top <code>k</code> candidates based on log probability. The top <code>k</code> candidates sorted in descending log probability is thus our list of recommended candidates. <h2 id="results-4"><a class="header" href="#results-4">Results</a></h2> Overall, the semantic ID approach was the best performing, and had significant gains over a fine-tuned dual encoder model. The scaling effects of larger models seemed much better with the DSI approach compared to dual encoders. In my experience, dual encoder models also don't seem to benefit much from scaling to larger models (minimal effect scaling from <code>20M</code> models to <code>500M</code> models.) The DSI approach is closer in performance with <code>hits@10</code> to dual encoder but had much better <code>hits@1</code> performance. <h2 id="concerns"><a class="header" href="#concerns">Concerns</a></h2> One finding was that shorter representations of the document (truncate at first <code>L</code> tokens) worked significantly better. The authors suggested that longer document lengths might make indexing memorization more difficult. This seems like a significant bottleneck to the paradigm, but future work might resolve this issue. One big shortcoming of this current method as-is is that the embedding representation used is an OTS T5 model. Subsequent papers will find that relevance-based fine-tuned embeddings perform much better at generating good semantic IDs to represent each document. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="rafailov-2023---direct-preference-optimization"><a class="header" href="#rafailov-2023---direct-preference-optimization">Rafailov 2023 - Direct Preference Optimization</a></h1> Here we trace the derivations from the <a href="https://arxiv.org/pdf/2305.18290.pdf">DPO paper</a>. Denote the model after SFT as πSFT(y∣x), i.e. the policy πSFT is a probability function for each pair of inputs and answers <code>(x, y)</code>. Naturally, we can use this policy to generate tokens by choosing the response with the highest probability (or approximate it in a greedy token-by-token manner). To perform RLHF, we first need to build a reward model. First, we prompt the SFT model to obtain pairs of answers (y1,y2)∼πSFT(y ∣ x). These samples are presented to human labellers who record their preferences in the form yw▹yl ∣ x, where yw wins and yl loses. These preferences are assumed to be generated from an underlying reward model r∗(x,y) which we do not have access to. We wish to learn this reward model. Since we only have access to pairwise preferences instead of the reward score, a common approach is to model the pairwise preferences using the <a href="papers/../misc/bradley-terry.html">Bradley-Terry</a> model. Specifically, we assume that the observed human preference decisions are related to the underlying human reward model in the following way: p∗(y1▹y2 ∣ x)=er∗(x, y1)+er∗(x, y2)er∗(x, y1)(1) Suppose we have a static dataset of comparisons D={ x(i),yw(i),yl(i) }i=1N sampled from the human preference model of p∗. We can create a reward model rϕ(x, y) and use the BT-model to express the negative log likelihood of D. Note that are we using the expression of <a href="papers/../misc/bradley-terry.html#bt-model-as-a-sigmoid-function">the BT-model as a sigmoid function</a>. With this NLL expression, we can optimize for rϕ using gradient descent to learn the reward model from D. Notice that the heart of Equation (2) is essentially just a difference in reward between the winning answer yw and losing answer yl. L(rϕ, D)=−E(x,yw,yl)∼D[ log σ(rϕ(x, yw)−rϕ(x, yl)) ](2) Note that rϕ is usually initialized using the SFT model πSFT, but with an additional linear layer over the final transformer layer to generate a scalar value as the reward. Having learned the reward model, we then need to use the learned reward function to fine-tune πSFT using reinforcement learning. Specifically, we set πref:=πSFT as the reference model, and initialize a new πθ model that we wish to train. Usually, πθ is also initialized as a copy of πSFT. The objective function we wish to maximize is: πθmaxE x∼D, y∼πθ(y∣x)[rϕ(x, y)]−β⋅DKL[πθ(y ∣ x) ∣∣ πref(y ∣ x)](3) Inspecting this objective, we see that we are trying to tune πθ such that it generates answers that maximize the learned reward function rϕ, while at the same time ensuring that we do not deviate too far from the original reference model. β is a hyperparameter controlling the degree of deviation. This penalty constraint serves to: <ol> <li>Ensure that we do not drift too far from the <code>(x, y)</code> distribution on which the reward model is accurate</li> <li>Ensure that we maintain generational diversity and not just collapse into a single high-reward answer for a given prompt</li> </ol> Objective (3) may not be directly optimized, because we need to generate y∼πθ(y ∣ x) at each step from the current policy (not sure if I fully understand this). Hence typically this is optimized using reinforcement learning using PPO. <h2 id="direct-preference-optimization"><a class="header" href="#direct-preference-optimization">Direct Preference Optimization</a></h2> The RLHF process in general is unstable, requires more memory / computation and requires tricks to make it work. Hence the authors of DPO set out to create an optimization procedure that: <ol> <li>Avoids fitting an explicit, standalone reward model</li> <li>Avoids using reinforcement learning</li> </ol> DPO starts off with the KL-constrained reward maximization objective from Equation (3) above. The first step is to show that the optimal policy for this objective is of the following form for an arbitrary reward model r(x,y): πr(y ∣ x)where partition fnZ(x)=Z(x)1πref(y ∣ x) exp (β1r(x, y))=y∑πref(y ∣ x) exp (β1r(x, y))(4) The derivation for Equation (4) is as follows. For a given reward function r(x,y): where Z(x)πmaxEx∼D,y∼π(y ∣ x)[r(x,y)]−β⋅DKL[π(y∣x) ∣∣ πref(y∣x)]=πmaxEx∼D,y∼π(y ∣ x)[r(x,y)−βlogπref(y∣x)π(y∣x)]=πminEx∼D,y∼π(y ∣ x)[logπref(y∣x)π(y∣x)−β1r(x,y)]=πminEx∼D,y∼π(y ∣ x)<svg xmlns="http://www.w3.org/2000/svg" width='0.667em' height='3.600em' viewBox='0 0 667 3600'><path d='M403 1759 V84 H666 V0 H319 V1759 v0 v1759 h347 v-84 H403z M403 1759 V0 H319 V1759 v0 v1759 h84z'/></svg>logZ(x)1πref(y∣x)exp(β1r(x,y))π(y∣x)−logZ(x)<svg xmlns="http://www.w3.org/2000/svg" width='0.667em' height='3.600em' viewBox='0 0 667 3600'><path d='M347 1759 V0 H0 V84 H263 V1759 v0 v1759 H0 v84 H347z M347 1759 V0 H263 V1759 v0 v1759 h84z'/></svg>=y∑πref(y ∣ x) exp (β1r(x, y)) For line 2 above, recall that DKL(P∣∣Q) is the expected value of logP(Y)−logQ(Y) if the random variable Y is drawn from P. Since the outer expectation is over draws from π, we can breakdown the KL-divergence by bringing the log difference into the expectation. Line 3 simply divides by −β and flips max to min. Line 4 uses logexp to bring the reward term into the denominator of the left term, then introduces an arbitrary Z(x). Note that the two Z(x) can be cancelled out if we brought them together, but we will be using them later on. Now let us define the optimal policy π∗(y∣x)=Z(x)1πref(y∣x)exp(β1r(x,y)). We will need to prove that π∗ is indeed optimal. Note that π∗ is a valid probability distribution as: <ul> <li>π∗(y∣x)≥0 ∀ y; and</li> <li>∑yπ∗(y∣x)=1, since the denominator is just the sum over y of the numerator</li> </ul> Since Z(x) is not a function of y, we can sub in π∗ and take logZ(x) out. The left term becomes a KL-divergence between π which we are optimizing over and the optimal policy π∗. πminEx∼D[Ey∼π(y ∣ x)[logπ∗(y∣x)π(y∣x)]−logZ(x)]=πminEx∼D[DKL(π(y∣x) ∣∣ π∗(y∣x))−logZ(x)] Finally, note that Z(x) does not depend on π, so we only need to consider the KL-divergence term. Gibb's inequality tells us that KL-divergence is minimized at 0 if and only if the two distributions are identical. This completes our derivation of (4) by showing that π∗ is indeed the optimal policy. Now that we have completed the derivation, let's consider what Equation (4) is saying. It tells us that we have an analytical solution for the policy πr that optimizes (3), and that it can be expressed in terms of πref (which we already have) and a given reward function r(x,y). Since we previously learned a reward model rϕ, we could simply plug that into (4) to get our optimal policy. Specifically, for a new input prompt x, we can compute πref(y∣x)exp(β1rϕ(x,y)) for all possible values of y and pick the best y. We can ignore Z(x) since it is fixed for a given prompt x. Intuitively, the new model scales the probability of high reward answers up with an exponential multiplier, and the degree of scaling is controlled by β. However, this is not computationally practical as we need to evaluate over a very large space of y∣x (i.e. all possible answers for a given prompt x). Hence, we want to find a form of the optimal policy which does not involve the partition function Z(x) nor the reward model r(x,y). We start by re-arranging Equation (4), taking log on both sides and re-arranging: r(x,y)=βlogπref(y ∣ x)πr(y ∣ x)+βlogZ(x)(5) Since Equation (5) holds for any arbitrary reward model, we can use the optimal (unknown) human reward model back in Equation (1), r:=r∗. Also, note that πr in Equation (5) refers to the optimal policy under r, so since we are using the optimal reward r∗, we can call this optimal policy π∗ as well. Now we plug Equation (5) back into Equation (1). p∗(y1▹y2 ∣ x)=exp(r∗(x, y1))+exp(r∗(x, y2))exp(r∗(x, y1))=exp(βlogπref(y1 ∣ x)π∗(y1 ∣ x)+βlogZ(x))+exp(βlogπref(y2 ∣ x)π∗(y2 ∣ x)+βlogZ(x))exp(βlogπref(y1 ∣ x)π∗(y1 ∣ x)+βlogZ(x))=1+exp(βlogπref(y2 ∣ x)π∗(y2 ∣ x)−βlogπref(y1 ∣ x)π∗(y1 ∣ x))1=σ(βlogπref(y2 ∣ x)π∗(y2 ∣ x)−βlogπref(y1 ∣ x)π∗(y1 ∣ x))(6) The derivation of the above is quite simple, we just need to note that it is of the form exp(x)/(exp(x)+exp(y)), and use the expression of <a href="papers/../misc/bradley-terry.html#bt-model-as-a-sigmoid-function">the BT-model as a sigmoid function</a>. The great thing is that the pesky partition function Z(x) cancels out because the BT-model simply ends up with the difference between two scores / rewards. Equation (6) looks simple but is quite remarkable when we compare it to Equation (1), because we now have the probability of the human preference data in terms of just the reference policy πref and the optimal policy π∗, without the reward model at all! Of course, the reward model is implicitly embedded in the equation. If we stare at Equation (6), we see the implicit reward function is r(x,y)=βlogπref(y ∣ x)π∗(y ∣x). The benefit of this new formulation is that we can write the maximum likelihood objective in terms of the optimal policy. We do not know the optimal policy, but we can parametrize it as πθ and use the maximum likelihood objective to train the model. Our new policy objective becomes: LDPO(πθ;πref)=E(x,yw,yl)∼D[logσ(βlogπref(yw ∣ x)πθ(yw ∣ x)−βlogπref(yl ∣ x)πθ(yl ∣ x))](7) <h2 id="application-to-dual-encoder-retrieval"><a class="header" href="#application-to-dual-encoder-retrieval">Application to Dual Encoder Retrieval</a></h2> The DPO framework offers a way to fine-tune a policy according to human preferences, whilst ensuring stability against a reference model. In the original formulation, π(y ∣ x) represents the probability of generative model generating y given an input prompt x. A related problem is that of fine-tuning embedding models for search retrieval, in what is known as the dual encoder framework. In this case, we also have preference data in the form of triplets <code>(query, positive_passage, negative_passage)</code>, and we wish to fine-tune embeddings such that the dot-product or cosine similarity between <code>(query, positive_passage)</code> is high whilst that of <code>(query, negative_passage)</code> is low. In this formulation, we could let the policy π(q, p) represent the normalized probability of a relevant match between query and passage. We can then borrow the framework of DPO to fine-tune our embeddings. The benefit of DPO compared to typical optimization objectives for dual encoder is the stability of the policy against the reference policy, which hopefully is a good form of regularization even when preference data is limited. Specifically, we define a dual encoder policy π(q, p), where q represents a query and p represents a passage, like so: <ol> <li>Encode q into a vector ∈Rd using a BERT model</li> <li>Encode p into a vector ∈Rd using a BERT model (can be same model as the query encoder)</li> <li>Take the dot product</li> <li>Run the dot product through a sigmoid layer to convert it into a probability</li> </ol> We can then apply this method to a reference encoder model and call it πref(q, p), and optimize another encoder model against this reference model. Will need to conduct experiments to see if this method offers gains against the typical dual encoder objectives, such as <a href="https://ar5iv.labs.arxiv.org/html/2004.04906">Karpuhkhin 2020</a>. <h2 id="references-5"><a class="header" href="#references-5">References</a></h2> <ul> <li><a href="https://arxiv.org/pdf/2203.02155.pdf">Ouyang 2022 - InstructGPT</a></li> <li><a href="https://arxiv.org/pdf/2305.18290.pdf">Rafailov 2023 - Direct Preference Optimization</a></li> <li><a href="https://huggingface.co/blog/rlhf">Huggingface RLHF Blog</a></li> </ul> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="blecher-2023---nougat"><a class="header" href="#blecher-2023---nougat">Blecher 2023 - Nougat</a></h1> <a href="https://arxiv.org/pdf/2308.13418">Blecher 2023 - Nougat: Neural Optical Understanding for Academic Documents</a>. Nougat is an end-to-end system that converts a scientific PDF into a sequence of tokens in markdown format in an auto-regressive way. Prior methods for Visual Document Understanding (VDU) usually rely on an external OCR service to generate intermediate outputs. In contrast, this method is end-to-end, and the text is generated directly from image embeddings in a decoder manner. Thus, the model is very simple and most of the work in this paper is in data preparation. <h2 id="model-2"><a class="header" href="#model-2">Model</a></h2> <ol> <li>Encoder. The encoder gets a variable size document image x∈R3×H0×W0 and applies crop / resizing to generate a fixed rectangle of size (H,W). Smaller images are white padded. The fixed size image can then be passed into a Swin Transformer to output a sequence of embedded patches z∈Rd×N where d is the latent dimension and n is the number of patches.</li> </ol> <h2 id=""><a class="header" href="#"></a></h2> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="dong-2023---mine-loss"><a class="header" href="#dong-2023---mine-loss">Dong 2023 - MINE Loss</a></h1> <a href="https://arxiv.org/abs/2312.08520">Revisiting Recommendation Loss Functions through Contrastive Learning</a> This paper compares several recommendation loss functions like BPR, CCL and introduce two new losses: InfoNCE+ and MINE+. <h2 id="setup-14"><a class="header" href="#setup-14">Setup</a></h2> Let U,I denote the user and item sets. We denote that: <ul> <li>Each user u∈U has interactions with Iu+ items, and has no interactions with the remaining I∖Iu+ set of items.</li> <li>On the item side, Ui+ denotes all users who interacted with item i</li> <li>We can also denote rui=1 if there was an interaction and rui=0 otherwise</li> </ul> Let the latent embeddings vu,vi represent user u and item i respectively. The similarity measure between them is then denoted y^ui=<vu,vi>. <h2 id="bpr-loss-1"><a class="header" href="#bpr-loss-1">BPR Loss</a></h2> The most widely used loss is Bayesian Personalized Ranking: LBPR=Eu Ei∼pu+j∼pi∑N−log σ(y^ui−y^uj)=Eu Ei∼pu+j∼pi∑N−log (1+exp(y^uj−y^ui)) Note that for each user, we take the expectation over the set of items relevant to him. We then sample N negatives from the overall item distribution (usually uniformly at random). <h2 id="softmax-loss"><a class="header" href="#softmax-loss">Softmax Loss</a></h2> One common approach is to model P(I=i∣U=u) as an extreme classification problem where I is a very large set. The probability may then be modeled as a softmax: P(I=i∣U=u)=∑j∈IevuTvjevuTvi In practice, it is infeasible to compute over the large item set, so we sample negative candidates for the denominator. The sampling is then corrected via importance weighting. <ul> <li><a href="https://static.googleusercontent.com/media/research.google.com/en//pubs/archive/45530.pdf">Covington 2016</a> for such an approach.</li> <li><a href="https://arxiv.org/abs/1412.2007">Jean 2014</a> for importance weighting correction</li> </ul> Using this approach, the loss may be formulated as: Lsoft=−Eulogi∈Iu+∑ey^ui+∑j∼pu−Ney^ujey^ui Note that pu− is a negative sampling distribution for each user u, and is typically implemented as p which is based on item popularity. <h2 id="contrastive-learning-loss-infonce"><a class="header" href="#contrastive-learning-loss-infonce">Contrastive Learning Loss (InfoNCE)</a></h2> InfoNCE loss looks very similar to the sampled softmax loss, although the motivation is different. The key idea is to pull similar points closer and push dissimilar points apart. InfoNCE loss is the most famous contrastive learning loss: Linfo=−Eui∈Iu+∑logey^ui+∑j∼pi−Ney^ujey^ui=−Eui∈Iu+∑<svg xmlns="http://www.w3.org/2000/svg" width='0.667em' height='3.600em' viewBox='0 0 667 3600'><path d='M403 1759 V84 H666 V0 H319 V1759 v0 v1759 h347 v-84 H403z M403 1759 V0 H319 V1759 v0 v1759 h84z'/></svg>y^ui−log<svg xmlns="http://www.w3.org/2000/svg" width='0.875em' height='3.600em' viewBox='0 0 875 3600'><path d='M863,9c0,-2,-2,-5,-6,-9c0,0,-17,0,-17,0c-12.7,0,-19.3,0.3,-20,1 c-5.3,5.3,-10.3,11,-15,17c-242.7,294.7,-395.3,682,-458,1162c-21.3,163.3,-33.3,349, -36,557 l0,84c0.2,6,0,26,0,60c2,159.3,10,310.7,24,454c53.3,528,210, 949.7,470,1265c4.7,6,9.7,11.7,15,17c0.7,0.7,7,1,19,1c0,0,18,0,18,0c4,-4,6,-7,6,-9 c0,-2.7,-3.3,-8.7,-10,-18c-135.3,-192.7,-235.5,-414.3,-300.5,-665c-65,-250.7,-102.5, -544.7,-112.5,-882c-2,-104,-3,-167,-3,-189 l0,-92c0,-162.7,5.7,-314,17,-454c20.7,-272,63.7,-513,129,-723c65.3, -210,155.3,-396.3,270,-559c6.7,-9.3,10,-15.3,10,-18z'/></svg>ey^ui+j∼pi−∑Ney^uj<svg xmlns="http://www.w3.org/2000/svg" width='0.875em' height='3.600em' viewBox='0 0 875 3600'><path d='M76,0c-16.7,0,-25,3,-25,9c0,2,2,6.3,6,13c21.3,28.7,42.3,60.3, 63,95c96.7,156.7,172.8,332.5,228.5,527.5c55.7,195,92.8,416.5,111.5,664.5 c11.3,139.3,17,290.7,17,454c0,28,1.7,43,3.3,45l0,9 c-3,4,-3.3,16.7,-3.3,38c0,162,-5.7,313.7,-17,455c-18.7,248,-55.8,469.3,-111.5,664 c-55.7,194.7,-131.8,370.3,-228.5,527c-20.7,34.7,-41.7,66.3,-63,95c-2,3.3,-4,7,-6,11 c0,7.3,5.7,11,17,11c0,0,11,0,11,0c9.3,0,14.3,-0.3,15,-1c5.3,-5.3,10.3,-11,15,-17 c242.7,-294.7,395.3,-681.7,458,-1161c21.3,-164.7,33.3,-350.7,36,-558 l0,-144c-2,-159.3,-10,-310.7,-24,-454c-53.3,-528,-210,-949.7, -470,-1265c-4.7,-6,-9.7,-11.7,-15,-17c-0.7,-0.7,-6.7,-1,-18,-1z'/></svg><svg xmlns="http://www.w3.org/2000/svg" width='0.667em' height='3.600em' viewBox='0 0 667 3600'><path d='M347 1759 V0 H0 V84 H263 V1759 v0 v1759 H0 v84 H347z M347 1759 V0 H263 V1759 v0 v1759 h84z'/></svg> Note that the only difference from the sampled softmax loss is that the log is inside the sum rather than outside. The InfoNCE loss has been shown to maximize the mutual information between user u and item i and minimize mutual information between unrelated pairs. <h2 id="empirical-exploration-of-infonce"><a class="header" href="#empirical-exploration-of-infonce">Empirical Exploration of InfoNCE+</a></h2> The authors propose an InfoNCE+, which is just adding some hyperparameters to InfoNCE and performing some empirical tuning of these hyperparameters. The InfoNCE+ proposes adding ϵ and λ: Linfo+N=−Eui∈Iu+∑(y^ui−λ⋅logN)=ϵ⋅ey^ui+j∼pu−∑Ney^uj Empirically, the authors find that setting ϵ=0 and λ=1.1 usually works best (tbh, the empirical evidence is not super convincing). <h2 id="theoretical-support-for-removing-positive-term-from-denominator"><a class="header" href="#theoretical-support-for-removing-positive-term-from-denominator">Theoretical Support for Removing Positive term from Denominator</a></h2> As we can see, setting ϵ=0 effectively removes the positive term ey^ui from the denominator of the loss. This makes intuitive sense as it would constrain y^ui from increasing which is what we want. This has theoretical backing as well, as explored in <a href="https://arxiv.org/abs/2110.06848">Decoupled Contrastive Learning - Yeh 2022</a>. The DCL paper also shows that removing the positive term from the denominator leads to more stable training and less hyperparameter sensitivity. The DCL loss is thus: LDCL=−Eui∈Iu+∑<svg xmlns="http://www.w3.org/2000/svg" width='0.875em' height='3.600em' viewBox='0 0 875 3600'><path d='M863,9c0,-2,-2,-5,-6,-9c0,0,-17,0,-17,0c-12.7,0,-19.3,0.3,-20,1 c-5.3,5.3,-10.3,11,-15,17c-242.7,294.7,-395.3,682,-458,1162c-21.3,163.3,-33.3,349, -36,557 l0,84c0.2,6,0,26,0,60c2,159.3,10,310.7,24,454c53.3,528,210, 949.7,470,1265c4.7,6,9.7,11.7,15,17c0.7,0.7,7,1,19,1c0,0,18,0,18,0c4,-4,6,-7,6,-9 c0,-2.7,-3.3,-8.7,-10,-18c-135.3,-192.7,-235.5,-414.3,-300.5,-665c-65,-250.7,-102.5, -544.7,-112.5,-882c-2,-104,-3,-167,-3,-189 l0,-92c0,-162.7,5.7,-314,17,-454c20.7,-272,63.7,-513,129,-723c65.3, -210,155.3,-396.3,270,-559c6.7,-9.3,10,-15.3,10,-18z'/></svg>y^ui−logj∼pu−∑Ney^uj<svg xmlns="http://www.w3.org/2000/svg" width='0.875em' height='3.600em' viewBox='0 0 875 3600'><path d='M76,0c-16.7,0,-25,3,-25,9c0,2,2,6.3,6,13c21.3,28.7,42.3,60.3, 63,95c96.7,156.7,172.8,332.5,228.5,527.5c55.7,195,92.8,416.5,111.5,664.5 c11.3,139.3,17,290.7,17,454c0,28,1.7,43,3.3,45l0,9 c-3,4,-3.3,16.7,-3.3,38c0,162,-5.7,313.7,-17,455c-18.7,248,-55.8,469.3,-111.5,664 c-55.7,194.7,-131.8,370.3,-228.5,527c-20.7,34.7,-41.7,66.3,-63,95c-2,3.3,-4,7,-6,11 c0,7.3,5.7,11,17,11c0,0,11,0,11,0c9.3,0,14.3,-0.3,15,-1c5.3,-5.3,10.3,-11,15,-17 c242.7,-294.7,395.3,-681.7,458,-1161c21.3,-164.7,33.3,-350.7,36,-558 l0,-144c-2,-159.3,-10,-310.7,-24,-454c-53.3,-528,-210,-949.7, -470,-1265c-4.7,-6,-9.7,-11.7,-15,-17c-0.7,-0.7,-6.7,-1,-18,-1z'/></svg> The authors also show that this "decoupling" is also justified from the <a href="https://arxiv.org/abs/1801.04062">Mutual Information Neural Estimator</a> perspective. Specifically, the MINE paper shows that we can estimate the true mutual information between each user u and item i by the following optimization problem: I(u,i)<svg xmlns="http://www.w3.org/2000/svg" width='100%' height='0.42em' viewBox='0 0 2364 420' preserveAspectRatio='none'><path d='M1181 0h2l1171 296c6 0 10 5 10 11l-2 23c-1 6-5 10 -11 10h-1L1182 67 15 340h-1c-6 0-10-4-11-10l-2-23c-1-6 4-11 10-11z'/></svg>=sup(vu;vi)Epu,i(y^ui)−logEpu⊗pi(ey^ui) Intuitively, we want to maximize the above equation over the similarity function parametrized by the embeddings vi,vu. <ul> <li>The first term takes an expectation of similarity scores over the joint user, item distribution where an interaction occurs (i.e. positive pairs).</li> <li>The second term takes an expectation of exponentiated similarity scores over the product measure of marginal user and item distributions (i.e. assuming independence between user and item distribution).</li> </ul> <h2 id="mine-loss"><a class="header" href="#mine-loss">MINE Loss</a></h2> The authors then say that a "simple" adaptation of the MINE problem to the recommendation setting is formalized as the MINE loss: Lmine=−EuEi∼pi+[y^ui−logEj∼pi(ey^uj)] Not too sure how this is derived from the above. They also add a hyperparameter λ to control the relative weightage of the positive and negative samples. This results in what they term as MINE+: Lmine+=−EuEi∼pi+[y^ui−λlogEj∼pi(ey^uj)] Based on some ablation studies, they find that λ=1.1−1.2 usually works best. The paper also offers some lower bound analysis and de-biasing of InfoNCE which I will not delve into for now. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="liu-2023---meaning-representations-from-trajectories"><a class="header" href="#liu-2023---meaning-representations-from-trajectories">Liu 2023 - Meaning Representations from Trajectories</a></h1> <a href="https://ar5iv.labs.arxiv.org/html/2310.18348">Paper Link</a> This paper demonstrates a way to compute sentence similarities using auto-regressive LLMs with a decoder-only architecture. In the encoder-decoder setting, characterized by BERT models, a sentence is typically represented by the embedding generated by the encoder at the final layer. Typically, either the <code>[CLS]</code> token embedding is taken, or the element-wise mean of the embeddings at the final layer is taken to represent the sentence. As <a href="papers/./reimers_2019.html">Reimers 2019</a> demonstrates, the default BERT trained using masked language modelling objective does not automatically have good embeddings suitable for Semantic Textual Similarity tasks. Hence the BERT model needs to be fine-tuned with some sentence similarity training task to generate good embeddings. This paper is interesting as it does not require any such fine-tuning. The main idea is to represent a sentence (sequence of tokens) by the distribution of text continuations (denoted as trajectories) generated by the auto-regressive LLM on the context sentence. In other words, as the paper elegantly puts it, we represent the meaning of linguistic terms by their usage. <h2 id="setup-15"><a class="header" href="#setup-15">Setup</a></h2> Let A denote a finite vocabulary of tokens and A∗ indicate the set of all variable-length finite sequences of tokens in A. A language model may be viewed as a map M(⋅∣⋅):A∗×A∗→[0,1] that associates a prompt sequence s∈A∗ and a possible continuation sequence t∈A∗ with a likelihood score M(t∣s)∈[0,1]. The paper chose inverse perplexity as the score: M(t=(a1 ... am)∣s):=i=1∏mPM(ai∣ s a1 ... ai−1)1/m Side note on inverse perplexity. We know that perplexity is as below. Hence the above score is indeed the inverse perplexity. Perplexity(Model M)=exp[−m1i=1∑mlogPM(ai∣ s a1 ... ai=1)]=exp[log(i=1∏mPM(ai∣ s a1 ... ai=1)−1/m)]=i=1∏mPM(ai∣ s a1 ... ai=1)−1/m Given two prompts u and v, we denote their semantic representation under model M as Mu and Mv. The semantic distance between them may be denoted d(Mu,Mv) where d is some distance function. The paper chose: d(Mu,Mv)=Et∼21(Mu+Mv)∣logMu(t)−logMv(t)∣=Et∼21(Mu+Mv)<svg xmlns="http://www.w3.org/2000/svg" width='0.333em' height='3.000em' viewBox='0 0 333 3000'><path d='M145 15 v585 v1800 v585 c2.667,10,9.667,15,21,15 c10,0,16.667,-5,20,-15 v-585 v-1800 v-585 c-2.667,-10,-9.667,-15,-21,-15 c-10,0,-16.667,5,-20,15z M188 15 H145 v585 v1800 v585 h43z'/></svg>m1i=1∑mlogPM(ai∣v,a<i)PM(ai∣u,a<i)<svg xmlns="http://www.w3.org/2000/svg" width='0.333em' height='3.000em' viewBox='0 0 333 3000'><path d='M145 15 v585 v1800 v585 c2.667,10,9.667,15,21,15 c10,0,16.667,-5,20,-15 v-585 v-1800 v-585 c-2.667,-10,-9.667,-15,-21,-15 c-10,0,-16.667,5,-20,15z M188 15 H145 v585 v1800 v585 h43z'/></svg> In other words, we sample a continuation trajectory t with equal probability from either prompt and generate a continuation sequence of length m, and compute the expected difference in log-likelihood of the generated continuation between the two models Mu and Mv. Since it is not feasible to integrate the above expression over all possible trajectories t, we sample n trajectories and compute an estimation to the above expected difference. <h2 id="results-5"><a class="header" href="#results-5">Results</a></h2> For the experiments, the authors set <code>n=20</code>, <code>m=20</code>, and sampling temperature λ<code>=1.0</code>. <ul> <li>Ablating on number of trajectories <code>n</code> shows that most of the performance is achieved with <code>n=20</code>, with marginal gains thereafter</li> <li>Ablating on sequence length <code>m</code> shows that most of the performance is achieved with <code>m=20</code>, with marginal gains thereafter. Interestingly, there is a big improvement in performance from <code>m=1</code> to <code>m=10</code>, showing that the distribution of the next token alone is not sufficiently discriminating as a semantic representation.</li> <li>Ablating on sampling temperature λ shows that <code>1.0</code> is optimal. Sampling from too diverse or too similar trajectories hurts the performance of this method.</li> </ul> The results show that this method is the best way to represent autoregressive models thus far, with even slightly better performance than encoder-based models like <code>T5</code> on the STS task. The results are not as good as models explicitly trained with unsupervised contrastive learning (see <a href="papers/./gao_2021.html">SimCSE</a>) for the STS task, but as the decoder model size is increased, the performance starts to reach similar levels. <h2 id="application-to-semantic-retrieval"><a class="header" href="#application-to-semantic-retrieval">Application to Semantic Retrieval</a></h2> Unfortunately, this method is not suitable for semantic retrieval, as the distance measure requires sampling from both query and document distributions (if we treat query q as one prompt and document d as the other prompt). Note that the setting is Semantic Textual Similarity where query and documents come from similar distributions. The authors create a proof-of-concept workaround by creating a fixed set of trajectories TD that is used instead of sampling t∼21(Mu+Mv), where they select n random documents from the dataset and create a sample trajectory from each document to form TD. We can then precompute Md for each document d on this set of trajectories TD. At inference time, we just need to compute Mq on the fixed set of trajectories, then compare the distance Mq against the pre-computed Md to get the distance scores. The authors show that this approximation obtains reasonable performance. However, I note that in typical semantic search settings, query and document usually come from different distributions where queries are short and documents are long. Hence, the continuations of a query and document might not be directly comparable. However, the beauty of this method is that since it is based on continuations, we can use prompting to adjust the behaviour of the LLM to somewhat align distributions. For example, we may on the document-side generate trajectories using a prompt like <code>Given document d, generate a query</code>. On the query-side, we may use something like <code>Given query q, generate other similar queries</code>. We can then compare trajectories generated in this fashion. Another issue with the original workaround is that of latency. Even with a fixed trajectory set, we need to compute M(q ∣ di) ∀ di∈TD at inference time, which is costly even for small number of trajectories. One way to deal with this issue is to have a hybrid setup: <ul> <li>On the <code>document</code> side, we use the high-precision decoder to generate a prospective set of queries from documents of size <code>n</code> (say <code>n=1,000</code>). Each document d is represented as a embedding of size <code>n</code> where each position records the likelihood of prospective query pi given document d using the decoder model.</li> <li>We use a separate fine-tuned encoder model to precompute embeddings for each of the prospective queries.</li> <li>On the <code>query</code> side, during inference, we use the same encoder model to embed the query and compute similarity against the <code>n</code> prospective queries. The similarity scores form a representation of size <code>n</code> for the query.</li> <li>We may then compute similarity between the size <code>n</code> query and size <code>n</code> document representations and perform retrieval with standard nearest neighbour vector databases.</li> </ul> In contrast to standard embedding vectors where each position encodes an unknown quality, these representation vectors have clear semantic meaning in each position (indicating the strength of relation to a particular prospective query). This setup also allows us to trade-off precision against recall by tweaking the sparsity of the representation vectors using score thresholding or top-k thresholding on either (or both) the query and document side. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="klenitskiy-2023---bert4rec-vs-sasrec"><a class="header" href="#klenitskiy-2023---bert4rec-vs-sasrec">Klenitskiy 2023 - BERT4Rec vs SASRec</a></h1> <a href="https://arxiv.org/abs/2309.07602">Turning Dross Into Gold Loss: is BERT4Rec really better than SASRec?</a> This is a short paper that argues the alleged performance gains of <a href="https://arxiv.org/abs/1904.06690">BERT4Rec</a> over <a href="papers/./kang_2018.html">SASRec</a> is not due to the masked cloze prediction task and bi-direction attention as the authors claim, but rather due to the loss function. BERT4Rec uses softmax cross entropy over the entire item catalog at each time step, whereas SASRec uses binary cross entropy against a single sampled negative at each time step. When the same softmax cross entropy is used for SASRec, it outperforms BERT4Rec consistently and trains faster. <h2 id="background-5"><a class="header" href="#background-5">Background</a></h2> Sequential recommendation is a popular approach currently to recommender systems, and in particular transformer models with self attention are the standard approach. <a href="papers/./kang_2018.html">SASRec</a> is the standard approach, where the task of sequential recommendation is treated as a causal modelling task where the self attention mechanism is only allowed to attend to previous time steps when making the prediction at time step t. BERT4Rec was proposed as an improvement to SASRec, and the claim was that introducing bi-directional attention (like BERT) and performing prediction on the cloze passage task (i.e. randomly masking x% of items) is able to lead to significant gains over SASRec. The argument is that the random masking is akin to data augmentation, as there are far more permutations of masked positions compared to just predicting the next item. The authors point out two misgivings they have with this interpretation, which is in line with my own intuitions: <ul> <li>BERT4Rec task is only weakly related to the final goal of sequential recommendations, whereas SASRec tasks for training and prediction are perfectly aligned (i.e. just predict the next item). This is akin to only using the BERT encoder (without the decoder) for a language modelling task, which is quite strange.</li> <li>BERT4Rec masks some items and only calculates losses for the subset of items, whereas SASRec computes losses for all items (except the last item) at once, getting more training signal from each training sequence</li> </ul> So we should expect SASRec to perform better and more efficiently than BERT4Rec. How then do we explain the performance discrepancy? The authors hypothesize that the performance difference is really due to the difference in loss functions between them, as explained below. <h2 id="setup-16"><a class="header" href="#setup-16">Setup</a></h2> Start with a set of users U and items I. Let each user u∈U be represented by a sequence of item interactions su={i1(u),i2(u),...,inu(u)}. Each sequential deep learning model may be abstracted as an encoder of input sequence su, and the encoded sequence be denoted Hu∈Rnu×d, where d is the latent dimension. To make predictions, given the full item embedding matrix E∈R∣I∣×d, we take: Ru=HuET∈Rnu×∣I∣ Then the t,i element of Ru may be denoted as rt,i(u) represents the predicted relevance of item i at time step t for user u. SASRec: Binary cross entropy loss. SASRec does not compute the full Ru prediction matrix. Instead, for each true positive item at each time step, it randomly samples one negative item and computes the predictions rt,it(u) and rt,−(u). Then the loss is: LBCE=−u∈U∑t=1∑nulog(σ(rt,it(u)))+log(1−σ(rt,−(u))) BERT4Rec: Softmax Cross Entropy. In contrast, BERT4Rec computes the full prediction matrix Ru for each user and computes the softmax over the entire item catalog for each masked item prediction. The cross entropy loss is thus: LCE=−u∈U∑t∈Tu∑log∑i∈Iexp(rt,i(u))exp(rt,it(u)) Note that for BERT4Rec, the inner summation is only over the time steps with masked items. If we were to translate this loss to SASRec, we would sum over all time steps. Sampled Softmax. Finally, it may not be computationally feasible to compute the softmax over the full item catalog. Hence the authors propose that for each user sequence in a batch, we sample N items a user has not interacted with and use the same set of negatives for each time step of a given sequence. Let INu− denote all items that user u has not interacted with. The loss is then: LCE−sampled=−u∈U∑t=1∑nulogexp(rt,it(u))+∑i∈INu−exp(rt,i(u))exp(rt,it(u)) <blockquote> Qn: This means that we have a different set of negatives for each user in a batch? Seems quite memory intensive. </blockquote> <h2 id="experiments-5"><a class="header" href="#experiments-5">Experiments</a></h2> The authors use the full sequence of item interactions for each user. The last (most recent) item is held out as the test set, and the second last item is chosen as the validation step. Models are trained with early stopping on the validation set. The authors note that the common practice of sampling negative items for metric computation is not a robust one, as it introduces randomness into the metrics. <blockquote> Note: I think the other reason sampling negatives is not robust is because it does not directly mirror the retrieval task, which requires choosing an item from the full catalogue. </blockquote> The results show that SASRec with <code>3,000</code> negatives is consistently the best model, beating BERT4Rec consistently. It also trains around 21 to 41 times faster than BERT4Rec. Hence the authors recommend sampled softmax SASRec as the de-facto standard instead of BERT4Rec. The authors do note that for the smaller datasets, SASRec may overfit relatively quickly (validation loss peaks and declines) and hence it is important to use early stopping. In contrast, BERT4Rec is more robust to overfitting and validation performance generally does not decline (it plateaus near the peak). <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="singh-2023---semantic-ids-for-recs"><a class="header" href="#singh-2023---semantic-ids-for-recs">Singh 2023 - Semantic IDs for Recs</a></h1> <a href="https://arxiv.org/abs/2306.08121">Better Generalization with Semantic IDs: A Case Study in Ranking for Recommendations</a> This paper proposes uses semantic IDs instead of hashing to random IDs to represent user or items. Semantic IDs mean that collisions to the same ID are semantically meaningful, and hence addresses the collision problem and also offers better generalization in cold start scenarios. The paper shows that this approach (i) retains memorization ability on par with random hashing and (ii) generalizes much better to unseen items and (iii) is much more computationally efficient as each user / item is represented simply by a <code>64-bit integer</code>. <h2 id="background-6"><a class="header" href="#background-6">Background</a></h2> In industry recommender systems, users or items are typically represented using random bloom hashing on a string ID (see <a href="papers/./weinberger_2009.html">Weinberger 2009</a>) into an embedding table. Such IDs have no semantic meaning, so collisions in this approach degrade performance. Nevertheless, it is empirically clear that the item-level memorization is important for good recommender performance. On the other extreme, one may choose to avoid IDs entirely and simply represent items using their content embedding. While it is clear that this obviates the collision problem and improves cold-start performance, multiple papers have shown that performance will be degraded overall due to inability to match the memorization ability of ID-based recommenders. The ranker used in experiments is from Google's video recommender. See <a href="https://arxiv.org/abs/2302.09178">Improving Training Stability for Multitask Ranking Models in Recommender Systems</a> and <a href="https://daiwk.github.io/assets/youtube-multitask.pdf">Recommending What Video to Watch Next: A Multitask Ranking System</a>. <h2 id="approach"><a class="header" href="#approach">Approach</a></h2> The approach pre-supposes that every item (video in this case) already has a good learned embedding. For Google's case, they have a video encoder that represents each YouTube video as a <code>2048</code> dimensional vector that captures the topicality of the video. The encoder takes both audio and visual features as input. The model training is described in <a href="https://openaccess.thecvf.com/content_CVPR_2020/papers/Lee_Large_Scale_Video_Representation_Learning_via_Relational_Graph_Clustering_CVPR_2020_paper.pdf">Large Scale Video Representation Learning via Relational Graph Clustering</a>. The approach has two stages: <ol> <li>Represent each item as a sequence of Semantic IDs. They employ a method called RQ-VAE to compress a dense content embedding into a few discrete tokens (represented as integers) which capture the semantic content.</li> <li>Train the ranking model with Semantic IDs. Once trained, the RQ-VAE model is frozen and used to generate semantic IDs for each item. We then train embeddings for each semantic ID along with the rest of the ranking model.</li> </ol> <h2 id="stage-1-rq-vae-for-semantic-ids"><a class="header" href="#stage-1-rq-vae-for-semantic-ids">Stage 1: RQ-VAE for Semantic IDs</a></h2> The idea of RQ-VAE is to iteratively find tokens that match the content embedding. Once a token is found, the next token will try to match the residual embedding and so on. Let x∈Rd be the content embedding. The algorithm is as follows: <ul> <li>Use encoder E to map the content embedding x∈RD to a latent vector z∈RD′</li> <li>The residual quantizer of L layers recursively quantizes z into semantic IDs <ul> <li>Each level l has a codebook Cl:={ekl}k=1K containing K vectors (<code>K=2048</code> in this paper)</li> <li>The residual rl for level l is used to find the nearest codebook vector el</li> <li>The ID associated with el is taken as the semantic ID</li> <li>We compute rl+1:=rl−el</li> <li>Thus we end up with a sequence of semantic IDs (e.g. <code>(1, 4, 6, 2)</code>)</li> </ul> </li> <li>A decoder D maps the quantized latent z^:=∑l=1Lel back to x^</li> </ul> The RQ-VAE model is trained with the following losses: <ul> <li>L:=Lrecon+Lrqvae</li> <li>Lrecon:=∣∣x−x^∣∣22 aims to reconstruct the content embedding x</li> <li>Lrqvae:=∑l=1Lβ⋅∣∣rl−sg[el]∣∣2+∣∣sg[rl]−el∣∣2 where sg is the stop-gradient operator (which disables gradient updates to the term in the operator). This is to encourage rl and the codebook vector el to move toward each other, so that we have a good codebook at the end.</li> </ul> <h2 id="stage-2-using-semantic-ids-for-ranking"><a class="header" href="#stage-2-using-semantic-ids-for-ranking">Stage 2: Using Semantic IDs for Ranking</a></h2> Now that each item v is represented by a sequence of semantic IDs (c1v,...,cLv), we can treat each ID as a token. The most intuitive thing is to treat each token as a subword and assign a unique embedding to each token (this is the <code>unigram</code> approach below). Unfortunately, this simplistic approach is not the most ideal. The experiments reveal that these semantic tokens behave more like characters in NLP, and it is important for performance to assign unique embeddings to sequences of tokens (i.e. create subwords out of the characters). There are two approaches that were experimented: <ul> <li>Create n-grams out of the IDs. Suppose an item has semantic IDs <code>(4, 6, 3, 2)</code>. A unigram approach would look up embeddings for each ID by itself. A bigram approach would look up embeddings for <code>(46, 63, 32)</code> and so on. Consequently, the embedding table size for a quantizer with <code>L</code> levels and <code>K</code> codes in each level is something like (L−N+1)×KN. This gets prohibitively expensive for larger <code>N</code>, so the experiments stop at bigram.</li> <li>Sentence piece based. As with natural language, most n-grams rarely occur (or do not occur) at all. Hence the n-gram approach is wasteful. The authors found that applying the <a href="papers/../nlp_course/intro.html#lecture-2-word-representation-and-text-classification">sentence piece approach</a> (I suppose using byte pair encoding?) is an effective way to learning the subwords to assign a unique embedding. The authors train the subword model on impressed items to a desired arbitrary vocab size. Note: this makes the approach less flexible, as we need to freeze both the quantizer and the set of subwords. However their experiments show that this is generally quite resilient.</li> </ul> <h2 id="results-6"><a class="header" href="#results-6">Results</a></h2> The experiments use CTR AUC as the target metric. The model is trained on a window of N days and predictions on day N + 1. Cold start performance refers to the performance on the subset of items that were newly introduced on day N + 1. Note that each user is represented as a sequence of items, including past item history and the current item. The item is represented by semantic IDs for itself. The main findings of the experiments: <ul> <li>Sentencepiece (SPM) approach > n-gram approaches for overall performance. As the vocab size is scaled up for the <code>SPM</code> approach, it outperforms both unigram and bigram easily.</li> <li>SPM > random hashing with the same vocab size for overall performance. It does not compromise on the memorization capacity of the model.</li> <li>All Semantic ID approaches > random hashing in the cold start scenario (as expected)</li> </ul> The authors also conducted a set of experiments where items are represented by their content embedding directly. Due to the large size of the embedding, they did not include user past history for these experiments. These showed that content embedding approach is inferior to random hashing, unless the number of layers is increased significantly. This suggests that a larger model is indeed able to memorize better just using the content embeddings, but at a significant computational cost. Hence this justifies the use of semantic IDs as a more efficient way to balance between memorization and cold start performance. <h2 id="takeaways-5"><a class="header" href="#takeaways-5">Takeaways</a></h2> Semantic IDs is a promising approach to balance memorization and cold start performance. However, it presupposes that we have good embeddings for each item. It also introduces significant engineering complexity in training and freezing a residual quantizer and a fixed subword vocabulary. Although the authors conduct experiments to show that this quantizer is quite robust to data distribution shift, there would probably come a time when we need to update the quantizer and the subword vocabulary and deal with the refreshing issue. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="yang-2023---opro"><a class="header" href="#yang-2023---opro">Yang 2023 - OPRO</a></h1> <a href="https://arxiv.org/abs/2309.03409">Large Language Models as Optimizers</a> This paper proposes an automated prompt optimization technique that resembles optimization. Compared to previous efforts, this paper makes use of an LLM's ability to understand patterns and what they call optimization trajectories to optimize a prompt. <h2 id="method-8"><a class="header" href="#method-8">Method</a></h2> We have access to a training set of <code><question, answer></code> pairs. The idea is simple: <ul> <li>Test out a prompt with the task description and evaluate the accuracy</li> <li>Append previously tested prompts and its accuracy in an ascending score order (keeping only the best <code>8</code> prompts in the trajectory)</li> <li>Ask the LLM to write a new prompt that is different from the old ones and achieve a high score</li> </ul> <h2 id="motivation"><a class="header" href="#motivation">Motivation</a></h2> The authors motivate this by showing that LLMs can solve simple linear regression and travelling salesman problems. For linear regression: <ul> <li>A true generating equation of y=wtruex+btrue+ϵ is chosen</li> <li><code>50</code> data points are generated with gaussian noise</li> <li>The goal is to get an LLM to guess what is the best <code>(w, b)</code> pair</li> <li>The algorithm loops through the following: <ul> <li>Start with <code>5</code> random <code>(w, b)</code> pairs, compute their mean-squared error</li> <li>Prompt the LLM with these pairs and tell it to suggest <code>8</code> new <code>(w, b)</code> pairs that further reduces the loss</li> <li>Evaluate the MSE of these new pairs and replace the pairs in the history</li> <li>Keep up to <code>20</code> pairs in history</li> </ul> </li> </ul> It is not obvious that LLMs can accomplish this task, as it requires some understanding of the underlying generating equation and how to nudge the <code>w, b</code> values. However, the paper shows that LLMs can generally solve this task and reach the global optimum, and that better models (e.g. <code>gpt-4</code> vs <code>gpt-3.5</code>) achieve the optimum in fewer number of steps. <h2 id="opro-for-prompt-optimization"><a class="header" href="#opro-for-prompt-optimization">OPRO for Prompt Optimization</a></h2> The general layout of a prompt is like so (note that the <code>Comment</code> headers are not part of the prompt): <pre><code class="language-markdown" style="white-space: pre-wrap;"> ### Comment: Meta Instruction for Optimization Trajectory ### I have some texts with their corresponding scores. The texts are arranged in ascending order based on their scores, where higher scores indicate better quality. ### Comment: Optimization Trajectory ### text: Let's figure it out! score: 61 ... more history ... ### Comment: Meta Instruction for Few Shot Exemplars ### The following exemplars show how to apply your text: you replace <INS> in each input with your text, then read the input and give an output. We say your output is wrong if your output is different from the given output, and we say your output is correct if they are the same. ### Comment: Few Shot Exemplars ### input: Q: Alannah, Beatrix, and Queen are preparing for the new school year and have been given books by their parents. Alannah has 20 more books than Beatrix. Queen has 1/5 times more books than Alannah. If Beatrix has 30 books, how many books do the three have together? A: <INS> output: 140 ... more exemplars ... ### Comment: Final Meta Instruction ### Write your new text that is different from the old ones and has a score as high as possible. Write the text in square brackets. </code></pre> Optimization is tested out on GSM8K (grade school math problems) and Big Bench Hard (logical reasoning questions). At each step, the above prompt is used <code>8</code> times to generate <code>8</code> different prompts. These prompts are evaluated on the training set. The best scoring <code>20</code> prompts are kept in the optimization trajectory at any point in time. Other parameters: <ul> <li><code>Temperature = 1</code> when generating prompts for diversity</li> <li><code>Temperature = 0</code> when evaluating each prompt on the training set with the scorer LLM for greedy decoding</li> </ul> The optimization trajectory looks like this, resembling actual optimization. <div class="table-wrapper"><table><thead><tr><th style="text-align: center"><img src="papers/../images/yang_2023_optimization_trajectory.png" alt="OPRO Optimization Trajectory" /></th></tr></thead><tbody> <tr><td style="text-align: center">OPRO Optimization Trajectory (Figure 1 from paper)</td></tr> </tbody></table> </div> <h2 id="notes"><a class="header" href="#notes">Notes</a></h2> The method is intuitive and compelling. The downside is that inclusion of both optimization trajectory and exemplars makes the method quite costly, although the optimization only needs to be done once to get the best prompt. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="rajput-2023---generative-retrieval"><a class="header" href="#rajput-2023---generative-retrieval">Rajput 2023 - Generative Retrieval</a></h1> <a href="https://arxiv.org/abs/2305.05065">Recommender Systems with Generative Retrieval</a> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="borisyuk-2024---gnn-at-linkedin"><a class="header" href="#borisyuk-2024---gnn-at-linkedin">Borisyuk 2024 - GNN at LinkedIn</a></h1> <a href="https://arxiv.org/abs/2402.11139v1">Borisyuk 2024 - LiGNN: Graph Neural Networks at LinkedIn</a> This is an applied paper documenting learning points from LinkedIn's latest Graph Neural Network (GNN) models, geared toward industry practitioners. The main architecture is based on <a href="https://ar5iv.labs.arxiv.org/html/1706.02216">GraphSAGE</a>. The graph is heterogenous, where nodes can represent a company or member or post etc., and edges can represent (i) engagement of a user with a post, (ii) affinity between a member and a creator, and (iii) whether a member or company has a particular attribute. Edges are weighted by the strength of the engagement, except attribute edges which all have a weight of <code>2.0</code>. The GNN's are trained as encoders for the various types of nodes, and the generated embeddings may be used for downstream link prediction tasks. They are trained using the GraphSAGE framework, which inductively generates the node embeddings using graph sampling and neighbourhood aggregation. The graph sampling is powered by the <a href="https://github.com/microsoft/DeepGNN">Microsoft DeepGNN Graph Engine</a>, which is run on a Kubernetes CPU cluster. <h2 id="temporal-gnn"><a class="header" href="#temporal-gnn">Temporal GNN</a></h2> One of the main innovations of this paper is to adapt the GraphSAGE framework to the temporal nature of a live recommender system, where items lose relevance very quickly with time and modelling the activity sequence is important. Firstly, the typical SAGE encoder is applied to the member and its neighbours (e.g. connections, affinities etc.). This generates an embedding of <code>d</code> dimensions. Multi-head attention is applied to generate a few of these embeddings, e.g. <code>H=4</code>. In addition, sampling of the member's past activities (e.g. N=100) is performed, and each of these activity nodes are also encoded. These are the embeddings labelled <code>A1, ..., A100</code> in the diagram below, note that the numbers correspond to chronological order. The input is now of dimension <code>H+N x d</code>. We can see how this now resembles a sequential input embedding (e.g. from text) which is typically fed into transformers and trained with a causal prediction task. Indeed, here we feed the sequential input into a transformer and get an output of <code>H+N x d</code> as well (referring to <code>H1, ..., H80</code> at the top). Positional encodings are added to the transformer so that the positions may be interpreted as a sequence. Finally, causal masking is used for training the transformer. This means that the first H tokens, no masking is applied, but for the last N positions, each token can only attend to the first H tokens and the activity tokens that preceded it. The temporal modelling contributes around 4-5% lift in AUC which is significant. <div class="table-wrapper"><table><thead><tr><th style="text-align: center"><img src="papers/../images/borisyuk_2024_temporal.png" alt="Temporal GNN Framework" /></th></tr></thead><tbody> <tr><td style="text-align: center">Temporal GNN Framework (Figure 4 from paper)</td></tr> </tbody></table> </div> <h2 id="graph-densification"><a class="header" href="#graph-densification">Graph Densification</a></h2> Since GNNs rely on neighbourhood aggregation, nodes with few interactions pose a cold start problem. To solve this problem, artificial edges are added between nodes with low out-degrees and nodes with high out-degrees based on content-based similarity. Specifically, their algorithm works like so: <ol> <li>Nodes with out-degree ≤ <code>30th percentile</code> designated as low-degree</li> <li>Nodes with out-degree ≥ <code>90th percentile</code> designated as high-degree</li> <li>Pre-embed all nodes with an LLM based on content</li> <li>For each low-degree node, find its approximate <code>top-50</code> neighbours amongst the high-degree set. Create an edge between each of them and the low-degree node with edge weight based on embedding similarity</li> </ol> It seems that graph densification adds a small (0.2%-0.5%) lift to AUC both on offline and online metrics. It probably also increases the diversity of recommendations. <h2 id="personalized-page-rank-graph-sampling"><a class="header" href="#personalized-page-rank-graph-sampling">Personalized Page Rank Graph Sampling</a></h2> Personalized Page Rank (PPR) is an adaptation of the well-known Page Rank algorithm and is widely used as a notion of node similarity. Specifically, given a source node <code>s</code> and a target node <code>t</code>, the PPR score π(s,t) represents the probability that a random walk originating at <code>s</code> will terminate at <code>t</code>. In training the GNN under the GraphSAGE framework, embeddings of neighbours to a node <code>s</code> are aggregated together to form the representation for <code>s</code>. A terse way of saying the same thing is that "GNNs propagate signals across graphs". Hence, how neighbours are sampled is crucial in determining the performance of the GNN. The simple way is Random Sampling, which simply chooses neighbours randomly amongst nodes connected to node <code>s</code>. This can done over a single-hop or multi-hops, and is efficient but not the most performant. The better way is PPR sampling, which chooses neighbours weighted by their PPR score, with the search space limited to nodes connected to <code>s</code> over <code>k</code> hops. This is slower but a better measure of neighbourhood. From their experiments, 2-hop PPR sampling contributes 2.1% validation AUC lift. Adding more hops beyond 2 hops contributes marginally to performance, hence 2-hops is chosen for efficiency. <h2 id="other-tips-and-tricks"><a class="header" href="#other-tips-and-tricks">Other Tips and Tricks</a></h2> The paper is full of practical tips on how to speed up training and improve stability. Here are some that stood out to me: Adaptive Graph Sampling. IO time is correlated with number of neighbours sampled. So they start training with a small number of neighbours and increases it only when validation metric starts to stall. It seems that the final number of neighbours is capped at <code>200</code>. Mixed Precision. Mixed precision (<code>float-16</code>) gave roughly 10% speed up in training. But they had to be careful to retain <code>float32</code> in the last layer otherwise it degraded the training process. Data Bound. One reflection point they made is that GNN training is generally data-bound rather than compute bound. So the focus on the paper was to make data loading and processing more efficient. For example, they implemented a shared memory queue, and use the <code>multiprocessing</code> package to have different processes pre-fetch data from the DeepGNN Graph Engine and pre-process it simultaneously. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="wang-2024---llm-for-pinterest-search"><a class="header" href="#wang-2024---llm-for-pinterest-search">Wang 2024 - LLM for Pinterest Search</a></h1> This paper explains Pinterest use of pre-trained LLMs to improve their search ranking model. They first fine-tune a Llama-8B model on human-labelled search relevance data, then used the model to label large amounts of unlabelled data. The unlabelled data was used to train a smaller MLP, which led to significant improvement in search performance. <h2 id="background-7"><a class="header" href="#background-7">Background</a></h2> Pinterest curates human-labelled dataset for search relevance. For each query, a pin is classified into five ordered relevance levels: <ul> <li>L5: Excellent / Highly Relevant</li> <li>L4: Good / Relevant</li> <li>L3: Complementary / Marginally Relevant</li> <li>L2: Poor / Irrelevant</li> <li>L1: Highly Irrelevant</li> </ul> <h2 id="teacher-model"><a class="header" href="#teacher-model">Teacher Model</a></h2> The teacher model takes the form of a standard cross encoder setup. The teacher model is either a BERT-based LLM or a decoder-style LLM. The text inputted to the LLM is of the following format: <pre><code>[CLS] Query [SEP] Pin Text </code></pre> The embedding of the <code>[CLS]</code> token is taken for BERT-based models, whilst the embedding of the final non-padding token is taken for decoder-based models (such as Llama). The embedding is passed through several fully-connected layers, and the final output dimension corresponds to the <code>5</code> relevance levels. The LLMs are fine-tuned during training by minimizing the pointwise multi-class cross entropy loss. To enhance the representation of the pin text for the teacher model, several features are concatenated together: <ul> <li>Pin titles and descriptions are used as-is</li> <li>Synthetic image captions are generated using an off-the-shelf model called BLIP</li> <li>High-engagement query tokens. Queries resulting in high engagement with a pin are aggregated over the past 2 years, and the top tokens are selected.</li> <li>User-curated Board Titles. Users curate their pins onto boards. The board titles are aggregated across users and top tokens are selected.</li> <li>Link Titles and Descriptions. A long click is defined as when a user clicks into a link to the webpage of a pin and stays <code>10 seconds</code> there. The incoming link url and description contains useful text which is mined.</li> </ul> A few off-the-shelf LLMs are tested as teacher models and the results reported below. <h2 id="student-model"><a class="header" href="#student-model">Student Model</a></h2> The teacher model is used to label large amounts of unlabelled search impressions data. Specifically, for each row of data, the teacher model generates a softmax probability distribution over the 5-point relevance labels. The student model is then taught to mimic the predicted probability distribution using cross-entropy loss. <blockquote> Note: The exact loss function for knowledge distillation is not covered in the paper. But the classical way is to perform knowledge distillation using KL-divergence loss. An implementation of the loss can be found <a href="https://nn.labml.ai/distillation/index.html">here</a>. The basic idea is to minimize KL-divergence between the teacher predicted probability distribution over labels and the student predicted probability distribution. </blockquote> The student model is a simple feed forward network on a diverse set of features. It seems like some feature engineering is performed to optimize performance. Features include: <ul> <li>Query-side embeddings. Pinterest has classifiers to categorize the query into a query interest type, shopping interest type etc. These categorical features are embedded using an embedding table. The SearchSage embedding for the query is also included.</li> <li>Pin-side embeddings. The PinSage pin embedding, SearchSage pin embedding etc. are included</li> <li>Query-pin interaction features. Standard query-pin match scores like BM25 score, % of query terms matched are also included as features</li> </ul> All features are presumably concatenated together, and passed into the MLP network. As mentioned above, this student network is trained using knowledge distillation from the teacher predictions above. <h2 id="results-7"><a class="header" href="#results-7">Results</a></h2> The teacher and student model are evaluated using human-annotated search relevance dataset. The train set for the teacher model is around 280k rows of human-annotated data, and the remaining 30k rows are used for evaluation. The teacher model is then used to run inference on 30 million rows of unlabelled data, and then used to train the student model. For the teacher model, scaling up the base model consistently produces better test performance (accuracy): <ul> <li>BERT-base: 53.5%</li> <li>T5-base: 56.9%</li> <li>DeBERTaV3-base: 58.0%</li> <li>Llama3-8b: 60.2%</li> </ul> Note that the Llama3-8b model was trained using qLora. The experiments also showed that: <ul> <li>Scaling up the amount of teacher-inferred data was crucial for improving student model performance</li> <li>The teacher LLM was able to successfully transfer learning from a purely English human-annotated dataset to multiple other languages, thus allowing the student model to learn multi-lingual behaviour.</li> <li>The inclusion of additional text features for the teacher model above helped to improve performance significantly.</li> </ul> <h2 id="takeaways-6"><a class="header" href="#takeaways-6">Takeaways</a></h2> This is a simple paper that shows the power of using LLM models to bootstrap learning from a smaller product-specific dataset to distill an effective student model. This is a useful paradigm because in search systems, we often have vast amounts of unlabelled impressions data that we can use a teacher model to run inference on. The paper also shows that scaling up the LLM is able to produce non-trivial increases in search relevance performance on the human-annotated dataset. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="solatorio-2024---gistembed"><a class="header" href="#solatorio-2024---gistembed">Solatorio 2024 - GISTEmbed</a></h1> <a href="https://arxiv.org/abs/2402.16829">GISTEmbed: Guided In-sample Selection of Training Negatives for Text Embedding Fine-tuning</a> This paper proposes a way to learn embeddings using contrastive learning. The main idea is to use a guide model to filter out false negatives from training for better data quality and performance. GISTEmbed models are currently topping the MTEB benchmark. <h2 id="implementation-4"><a class="header" href="#implementation-4">Implementation</a></h2> The method has an <a href="https://github.com/UKPLab/sentence-transformers/blob/master/sentence_transformers/losses/GISTEmbedLoss.py">implementation</a> in the SentenceTransformer package. We will walk through the implementation here. Firstly, the loss is initialized with both a <code>model</code> to train and a <code>guide</code> model, both of which are <code>SentenceTransformers</code>. In the <code>forward</code> step: <ul> <li><code>sentence_features</code> is the input which is a list of <code>dict[str, Tensor]</code> <ul> <li>The list is length <code>2</code> if we have anchor, positive</li> <li>The list is length <code>3</code> if we have anchor, positive, negative</li> </ul> </li> <li>The <code>sentence_features</code> is passed into both the <code>model</code> and the <code>guide</code> to get the respective embeddings <ul> <li>For <code>guide</code>, we may need to re-tokenize it if the tokenizer of the <code>guide</code> differs from <code>model</code></li> <li>This is done by using <code>batch_decode</code> and then <code>tokenize</code> again</li> </ul> </li> <li>Now we have <code>anchor</code>, <code>positive</code> and <code>negative</code> embeddings, each of shape <code>(batch_size, embed_dim)</code></li> <li>The <code>sim_matrix</code> is used to compute pairwise cosine similarities: <ul> <li>The code is simply <code>torch.nn.CosineSimilarity(dim=-1)(embed1.unsqueeze(1), embed2.unsqueeze(0))</code></li> <li><code>embed1</code> becomes shape <code>(batch_size, 1, embed_dim)</code> and embed2 becomes shape <code>(1, batch_size, embed_dim)</code></li> <li>The similarity is compared at dimension <code>-1</code></li> <li>Broadcasting ensures that the comparison is done pairwise, such that the result is of shape <code>(batch_size, batch_size)</code></li> <li>This is a common way to do pairwise similarity</li> </ul> </li> <li>Now we obtain the pairwise similarity matrices: <ul> <li><code>ap_sim</code>, <code>aa_sim</code>, <code>pp_sim</code>, <code>an_sim</code></li> <li><code>guided_ap_sim</code>, <code>guided_aa_sim</code>, <code>guided_pp_sim</code>, <code>guided_an_sim</code></li> </ul> </li> <li>The anchor positive similarity threshold is used to filter away potential false negatives <ul> <li>This is simply the <code>guided_ap_sim.diagonal()</code> which corresponds to the similarity between the anchor and positive in each row</li> <li>Note that they use the guide model for determining the threshold</li> <li>This threshold is called <code>guided_sim</code></li> </ul> </li> <li><code>mask_false_negatives</code> is used to suppress false negatives <ul> <li>Using the <code>absolute</code> strategy, cases where the <code>guided_sim_mat > guided_sim - self.margin</code> will be suppressed (set to <code>torch.inf</code>)</li> <li>The idea is that negatives should not have a higher similarity than the threshold, otherwise there is a higher probability they are false negatives</li> <li>This function is applied to <code>ap_sim</code>, <code>aa_sim</code>, <code>pp_sim</code> and <code>an_sim</code> to mask false negatives</li> </ul> </li> <li>Finally, we have <code>scores</code> <ul> <li><code>scores = torch.cat([ap_sim, aa_sim, pp_sim, an_sim], dim=1) / self.temperature</code></li> <li>This is of shape <code>(batch_size, 4*batch_size)</code></li> </ul> </li> <li>We create <code>labels</code> which is <code>torch.arange(batch_size)</code> <ul> <li>This is because the correct label is in the diagonal of <code>scores</code> matrix</li> </ul> </li> <li>Finally the loss is computed via <code>torch.nn.CrossEntropyLoss()(scores, labels)</code> <ul> <li>Each row is considered as a classification task where the label corresponds to the column position where the correct class is found</li> <li>The log softmax loss is then computed per row and averaged across rows</li> </ul> </li> </ul> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="sanjabi-2025---360brew"><a class="header" href="#sanjabi-2025---360brew">Sanjabi 2025 - 360Brew</a></h1> <a href="https://arxiv.org/pdf/2501.16450">360Brew: A Decoder-only Foundation Model for Personalized Ranking and Recommendation</a> 360Brew is a 150B foundational LLM model used to centralize many LinkedIn models into one model. For V1.0, 360Brew focused on replacing ranking tasks since they are less bounded by computational constraints in practice compared to retrieval. <h2 id="current-paradigm"><a class="header" href="#current-paradigm">Current Paradigm</a></h2> The current paradigm relies on bloom embedding of ID based features for both members and items. These large embedding tables are supplemented with other engineered features such as item attributes. Some challenges with the current paradigm are: <ul> <li>Cold start: ID based features cannot handle cold start, and thus a lot of work is done to learn good content-based representations</li> <li>Feature interactions: Previously,</li> </ul> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="zhang-2025---qwen3-embedding"><a class="header" href="#zhang-2025---qwen3-embedding">Zhang 2025 - Qwen3 Embedding</a></h1> <a href="https://arxiv.org/abs/2506.05176">Qwen3 Embedding: Advancing Text Embedding and Reranking Through Foundation Models</a> This short paper introduces the Qwen3 Embedding and reranker series, which are the strongest open source models currently for such tasks. The Qwen3 foundation decoder LLM serves as the backbone for fine-tuning and also is used to generate high quality synthetic training data. Note that Qwen3 is multi-lingual and are publicly available under the Apache 2.0 license, which means it can be used commercially. <h2 id="characteristics"><a class="header" href="#characteristics">Characteristics</a></h2> The embedding and reranking models come in 3 sizes: <code>0.6B</code>, <code>4B</code> and <code>8B</code>. <ul> <li><code>0.6B</code>: 28 layers, embedding dimension of <code>1024</code> for the embedder</li> <li><code>4B</code>: 36 layers, embedding dimension of <code>2560</code> for the embedder</li> <li><code>8B</code>: 36 layers, embedding dimension of <code>4096</code> for the embedder</li> </ul> Since the <code>4B</code> and <code>8B</code> have same number of layers, presumably the <code>8B</code> model has larger hidden sizes. All the models have <code>32K</code> sequence length limit, and are instruction aware, meaning that we can adjust the instruction at the start of the prompt to adjust the behaviour of the embedder or reranker. For the embedding models, there is also MRL support (Matryoshka Representation Learning), meaning that we can use custom dimensions from the embeddings. <h2 id="embedder"><a class="header" href="#embedder">Embedder</a></h2> The text embeddings are obtained by appending an <code>[EOS]</code> token at the end of every input sequence. The final embedding is derived from the hidden state of the last layer corresponding to this <code>[EOS]</code> token. Input format for queries or Documents is as follows: <pre><code>{Instruction} {Query or Document}<|endoftext|> </code></pre> The contrastive loss based on InfoNCE is used for training the embedder. Specifically, given a batch of N training instances, the loss is defined as: Lembedding=−N1i∑NlogZies(qi,di+)/τ s is cosine similarity function, τ is temperature and Zi is the normalization factor which includes the positive pair <code>+</code> various negative pairs: Zi=es(qi,di+)/τ+k∑Kmikes(qi,di,k−)/τ+j=i∑Nmijes(qi,qj)/τ+j=i∑mijes(di+,dj)/τ Comment on the above normalization factor: <ul> <li>The second term is the similarity between each anchor query and K hard negatives di,k− per query. Note that as it is written, only the hard negatives in the same row are used as negatives for each anchor query, but in theory we could use all negatives in the mini-batch.</li> <li>The third term is the similarity between pairs of queries. The assumption is that randomly selected queries should be unrelated to each other.</li> <li>The last term is the similarity between the positive document in each row (i.e. di+) and all other documents dj (including hard negatives).</li> </ul> The mij and mik are mask factors designed to reduce impact of false negatives in the normalization factor Zi. Specifically, given an anchor query or document i and a potential negative query or document j: mij={01if sij>s(qi,di+)+0.1 or dj==di+,otherwise This means that for each row, we use the similarity score between the query qi and di+ as a dynamic threshold to filter out false negatives. For any term in Zi which has too high similarity exceeding this threshold (plus a small margin), we reject the false negative and mask it out. Note that this approach is reminiscent of triplet loss semi-hard masking or the GISTEmbed loss. <h2 id="reranker"><a class="header" href="#reranker">Reranker</a></h2> The reranker is simpler, and training remains in the text paradigm. Specifically, the authors use the LLM chat template to incorporate instruction and frame the reranking task as a <code>yes</code> or <code>no</code> question: <pre><code><|im_start|>system Judge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be "yes" or "no".<|im_end|> <|im_start|>user <Instruct>: {Instruction} <Query>: {Query} <Document>: {Document}<|im_end|> <|im_start|>assistant <think>\n\n</think>\n\n </code></pre> Instead of fitting a classifier head, no change is made to the architecture. The reranking score is computed as the likelihood ratio of the next token being <code>yes</code> or <code>no</code>: score(q,d)=eP(yes∣I,q,d)+eP(no∣I,q,d)eP(yes∣I,q,d) The task then reduces to a supervised fine-tuning task, where the label is either <code>yes</code> for positives or <code>no</code> for negatives. The loss is simply the log probability of the correct label for each row (<code>yes</code> or <code>no</code>). Lreranking=−logp(l∣q,d) <h2 id="multi-stage-training"><a class="header" href="#multi-stage-training">Multi-stage Training</a></h2> The multi-stage training has emerged as a common practice for training text embedding models. The 3 stages used are as follows: <ul> <li>Stage 1: Large scale synthetic data. The Qwen 32B model was used to synthesize training pairs of data across many tasks, such as retrieval, classification, semantic textual similarity. <ul> <li>To create diversity, a document is taken from the Qwen3 training corpus, and top 5 similar documents are retrieved</li> <li>Qwen3 is presented with these documents and a user persona to generate a potential query</li> <li>Qwen3 is also instructed to vary the query type, length, difficulty and language for each query</li> <li>150 million query - document pairs are generated this way</li> </ul> </li> <li>Stage 2: High quality synthetic data <ul> <li>The 150 million pairs in Stage 1 are filtered down to 12 million high quality pairs</li> <li>Specifically, only query - document pairs with cosine similarity greater than 0.7 are selected</li> </ul> </li> <li>Stage 3: Model merging <ul> <li>Model merging based on Spherical Linear Interpolation is used, which merges multiple model checkpoints saved during the fine tuning process.</li> </ul> </li> </ul> Note that all 3 stages were used for the embedder, but stage 1 was omitted for the reranker as it did not help. The ablation studies show that all 3 stages are crucial for final performance of the <code>0.6B</code> embedding model. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="chlon-2025---llms-are-bayesian-in-expectation"><a class="header" href="#chlon-2025---llms-are-bayesian-in-expectation">Chlon 2025 - LLMs are Bayesian in Expectation</a></h1> <a href="https://arxiv.org/abs/2507.11768">LLMs are Bayesian, in Expectation, not in Realization</a> This paper uses information theoretic lens to analyze the effect of positional encoding on permutation qualities of LLMs and propose an algorithm to derive the optimal chain of thought length. <h2 id="summary"><a class="header" href="#summary">Summary</a></h2> In context learning allows LLMs to adapt to new tasks using only a few examples at inference time. A theoretical framework for interpreting ICL is through the lens of bayesian inference (see <a href="https://arxiv.org/abs/2111.02080">Xie 2022 - ICL as Implicit Bayesian Inference</a>). It proposes that transformers implicitly perform posterior updates over a latent concept variable, with the pretraining distribution encoding a prior over possible tasks. This perspective has been challenged by <a href="https://arxiv.org/abs/2406.00793">Falck 2024 - Is ICL in LLMs bayesian?</a>, which demonstrates empirically that transformer based language models systematically violate the martingale property. Specifically, for exchangeable data where the order of observations carries no information, bayesian posterior predictive distributions must satisfy: E[f(Xn+1)∣X1,...,Xn]=E[f(Xn+1)∣Xπ(1),...,Xπ(n)] For any permutation π and bounded function f. The experiments show that LLMs like GPT-3.5 consistently violate this property under input permutation. This paper observes that while bayesian inference assumes exchangeable data, position encodings fundamentally break the symmetry. This is formalized using two complexity measures: <ul> <li>Kolmogorov complexity K(X) of a sequence, which is permutation invariant for exchangeable data</li> <li>Conditional complexity K(X∣π) given a specific ordering π</li> </ul> It then shows that transformers with position encoding minimizes: Eπ∼U(Sn)[K(X∣π)]=K(X)+I(X;π) Where: <ul> <li>U(Sn) denotes the uniform distribution over permutations consistent with sufficient statistics of the data <ul> <li>For iid data, we can just sample uniformly over all permutations</li> </ul> </li> <li>I(X;π) is the mutual information between sequences and their orderings</li> </ul> Note that this is a well known theorem (the kolomogorov version of Shannon's information identity) applied to this context. <h2 id="notations"><a class="header" href="#notations">Notations</a></h2> <ul> <li>X=(x1,...,xn) denotes a sequence of observations</li> <li>Sn=∑i=1nxi is the sufficient statistic for bernoulli sequences</li> <li>π∈Sn represents a permutation of n elements</li> <li>Tθ denotes a transformer with parameters θ</li> <li>K(X) is the kolmogorv complexity of sequence X</li> <li>H(p)=−plogp−(1−p)log(1−p) is the binary entropy function</li> <li>k denotes the number of chain of thought tokens</li> <li>ϵ denotes the target error tolerance</li> </ul> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="huang-2025---llm-jepa"><a class="header" href="#huang-2025---llm-jepa">Huang 2025 - LLM-JEPA</a></h1> This paper proposes to add a JEPA loss to regular supervised fine-tuning of LLMs. The JEPA loss tries to minimize distance between the encoding of two pairs of texts that should share the same representation (e.g. <code>text</code>: description of a regex, <code>code</code>: code representing the regex). <h2 id="premise"><a class="header" href="#premise">Premise</a></h2> The next token prediction task of LLMs is not very amenable to JEPA-style training. However, there are subtasks that are well suited for JEPA - tasks involving two representations of the same reality. For example, in a "Natural Language to Regex" task, one may view the natural language instruction of <code>lines containing the string dog two or more times</code> and the regex <code>.*(dog)(2,).*</code> as two views of the same thing. Normal SFT would just try to predict the regex <code>code</code>, token-by-token, from the natural language instruction (<code>text</code>). This paper advocates for adding a JEPA-style loss which tries to align the latent representation of the prediction given the input <code>text</code> with the latent representation of the <code>code</code>. <h2 id="method-9"><a class="header" href="#method-9">Method</a></h2> Specifically, we can express the usual SFT objective of LLM training as: LLLM(Text1:L−1,TextL)=CrossEntropy(Classifier(Enc(Text1:L−1,TextL))) That is, the loss is the cross entropy of the next token predicted by our LLM classifier against the tokens of the label. The SFT loss is usually the sum of log losses over the label tokens. The LLM-JEPA objective adds an additional loss term to the usual SFT loss. The new loss aims to align the predictor representation of the input <code>text</code> with the encoder presentation of the <code>code</code>. LLLM−JEPA=l=2∑LLLLM(Text1:l−1,Textl)+λ⋅d(Pred(Enc(Text)),Enc(Code)) Some details: <ul> <li>λ is a hyperparameter balancing the contribution of the two terms</li> <li>d is the cosine similarity loss in this paper, i.e. 1−CosSim(⋅)</li> <li>Encoder. The encoding for text and code is taken as the <code>hidden_state</code> of the last token from the last layer. <ul> <li>To avoid two forward passes, this paper concatenates <code>text</code> and <code>code</code> to do a single forward pass and extracts the relatively embeddings</li> <li>The <code>attention_mask</code> is carefully modified so that <code>text</code> and <code>code</code> never attend to each other</li> </ul> </li> <li>Predictor. Instead of adding more layers, we insert one or more special <code>[PRED]</code> tokens (<code>k</code> tokens) to the end of the concatenated sequence. <ul> <li>These tokens are only allowed to attend to the <code>text</code> segment (?)</li> <li>We may view these as enabling additional non-linear computation to process the encoding of the <code>text</code> into a latent prediction of the <code>code</code></li> <li>It is interesting because in theory, the model could learn to generate a good latent prediction on these <code>[PRED]</code> tokens without affecting its next token prediction. So the fact that this separate task helps the NTP task is non-trivial.</li> </ul> </li> <li>In sum, there are two forward passes during training Note: To check on this part: <ul> <li>First pass: Over the input + label sequence (with <code>[PRED]</code> tokens at the end of the input) with standard causal masking <ul> <li>Allows us to compute the usual SFT loss: LLLM</li> <li>Allows us to extract Pred(Enc(Text))</li> </ul> </li> <li>Second pass: Over the input + label sequence with special blocked masking <ul> <li>Allows us to extract Enc(Code)</li> </ul> </li> </ul> </li> </ul> Finally, because of the increased computational cost, the authors considered dropout to reduce compute and also improve robustness. Specifically, we randomly choose forward passes where only LLLM is computed. They found that we can have high levels of dropout, e.g. LD=0.5 while improving performance and reducing compute cost. <h2 id="findings-1"><a class="header" href="#findings-1">Findings</a></h2> The paper spends quite a bit of time exploring the benefits of the JEPA loss. Here we summarize: <ul> <li>LLM-JEPA aligns the encoding of <code>text</code> and <code>code</code>. Specifically, it results in a representation where Enc(Code) is almost a linear transformation from Enc(Text). Without the JEPA loss, it is very clearly not the case. Using t-SNE, they also show two nice clusters with JEPA loss.</li> <li>LLM-JEPA improves fine-tuning accuracy. The latent prediction task improves the NTP prediction capability of the LLM as measured by accuracy on downstream tasks. The gains are especially large on the regex synthesis task.</li> </ul> <h2 id="hyperparameters"><a class="header" href="#hyperparameters">Hyperparameters</a></h2> There are two main hyperparameters to tune, which is not ideal: <ul> <li>λ controls the relative weight between LLLM and LJEPA. The paper searches between λ=1,...,4.</li> <li>k controls the number of predictor tokens (representing the computational capacity allocated to the predictor). The paper searches between k=0,...,4</li> </ul> <h2 id="code"><a class="header" href="#code">Code</a></h2> Here we do a deep dive into the <a href="https://github.com/rbalestr-lab/llm-jepa">code repository for LLM-JEPA</a>. <h3 id="datasets"><a class="header" href="#datasets">datasets/</a></h3> The <code>datasets/</code> folder comprises the <code>train</code> and <code>test</code> splits for different benchmarks stored in <code>jsonl</code> files. For example, one line of <code>gsm8k</code> looks like below. The aim is to produce the correct answer (exact match) given the system and user prompt. <pre><code class="language-json">{"messages": [ {"role": "system", "content": "Answer the math question, show steps."}, {"role": "user", "content": "James decides to bulk up. He weighs 120 kg and gains 20% of his body weight in muscle and 1 quarter that much in fat. How much does he weigh now?" }, {"role": "assistant", "content": "He gains 120*.2=120*.2=24 24 kg in muscles So he gains 24/4=24/4=6 6 kg of fat That means he gains a total of 24+6=24+6=30 30 kg of bodyweight So his new body weight is 120+30=120+30=150 150 kg #### 150"} ] } </code></pre> Note that <code>messages</code> has three components: <code>system</code>, <code>user</code>, <code>assistant</code>. <h3 id="finetunepy"><a class="header" href="#finetunepy">finetune.py</a></h3> The first part of the script handles dataset loading. load_and_prepare_dataset is the main function for data loading and tokenization. <ul> <li><code>data_file</code>: the <code>jsonl</code> file to load</li> <li><code>tokenizer</code>: the huggingface tokenizer we are using</li> <li><code>model_name: str</code>: the huggingface model we are using</li> <li><code>predictors</code>: the number of prediction tokens we will append to the end</li> <li><code>regular: bool</code>: If True, do normal LLM SFT, otherwise JEPA. regular is much simpler</li> <li><code>train_all: bool</code>: If True, compute loss on all tokens (not just <code>assistant</code>)</li> <li><code>plain: bool</code>: If True, do not <code>apply_chat_template</code>, instead format manually <ul> <li>Used when chat template is not available for a particular model</li> </ul> </li> <li><code>front_pred: bool</code>: If True, put predictor tokens at the start of sequence (not effective).</li> <li><code>reverse_pred: bool</code>: If True, swap user and assistant (predict <code>user</code> from <code>assistant</code>)</li> </ul> Step by step: <ul> <li>First we use <code>load_dataset</code> from huggingface <code>datasets</code> <ul> <li>Always load the <code>train</code> split as train and test are stored in different files</li> </ul> </li> <li></li> </ul> <pre><code>Input JSONL: ┌─────────────────────────────────────────┐ │ {"messages": [system, user, assistant]} │ └─────────────────────────────────────────┘ │ ▼ ┌──────────┼──────────┐ │ │ │ ▼ ▼ ▼ ┌────────┐ ┌────────┐ ┌──────────┐ │ Full │ │ User │ │Assistant │ │ Conv │ │ Only │ │ Only │ └────────┘ └────────┘ └──────────┘ │ │ │ ▼ ▼ ▼ LM Loss User Repr Target Repr └─────┬────┘ ▼ JEPA Loss (cos similarity) </code></pre> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="penha-2025---joint-task-semantic-id"><a class="header" href="#penha-2025---joint-task-semantic-id">Penha 2025 - Joint-task Semantic ID</a></h1> <a href="https://arxiv.org/abs/2508.10478">Semantic IDs for Joint Generative Search and Recommendation</a> This paper addresses the task of constructing semantic IDs that are useful for both search and recommendations. The main empirical finding is that constructing semantic IDs in a task specific way will necessarily degrade performance for the other task. <h2 id="methods-1"><a class="header" href="#methods-1">Methods</a></h2> The baseline / single-task methods considered for constructing semantic IDs: <ul> <li>Content-based. An off-the-shelf embedder is used to embed textual content of items. This is the approach used in DSI, TIGER. Specifically, <code>all-mpnet-base-v2</code> was used on the concatenated metadata of each item. The embedding is then discretized.</li> <li>Search-tuned IDs. Starting from <code>all-mpnet-base-v2</code>, the model is fine-tuned on search data with in-batch random negatives (i.e. <code>MultipleNegativesRankingLoss</code> from <code>sentence-transformers</code>). The search data comprises of (<code>search_query</code>, <code>relevant_item_metadata</code>) pairs. The fine-tuned embedding is then discretized.</li> <li>Recommendation-tuned IDs. The Efficient Neural Matrix Factorization (EMNF) method from TokenRec is used to create collaborative filtering based embeddings vrec.</li> </ul> <h2 id="multi-task-methods"><a class="header" href="#multi-task-methods">Multi-Task Methods</a></h2> A few methods were explored for how to combine multi-task signals: <ul> <li>Separate. This means that each task has its own set of item IDs. At inference time, search prompts can only output search tokens, and recommendation prompts can only output rec tokens.</li> <li>Fused Concat. Each embedding vsearch and vrec are individually l2-normalized and concatenated together.</li> <li>Fused SVD. Each embedding is l2-normalized but we further dimension reduce the embedding with higher dimensional space using truncated SVD. The two embeddings are then element-wise added together.</li> <li>Multi-task. A bi-encoder is trained on both supervision signals: <ul> <li><code>(query, item)</code> pairs from search data and <code>(item_a, item_b)</code> pairs from interaction data.</li> </ul> </li> </ul> <h2 id="semantic-id-learning-methods"><a class="header" href="#semantic-id-learning-methods">Semantic-id Learning Methods</a></h2> The methods considered for learning semantic ID: <ul> <li>RQ-kmeans. This is basically hierarchical k-means on the residuals at each level, implemented using FAISS residual quantizer.</li> <li>RQ-VAE. Implemented using <code>vector-quantize-pytorch</code></li> <li>MiniBatchDictionaryEncoding. From <code>sklearn</code> library</li> <li>ResidualLFQ from <code>vector-quantize-pytorch</code></li> </ul> <h2 id="data--metrics"><a class="header" href="#data--metrics">Data / Metrics</a></h2> A search and recommendation dataset is built from <code>MovieLens25M</code>: <ul> <li><code>62k</code> movies</li> <li><code>1.2M</code> user-item interactions (last item per user used for test)</li> <li><code>20</code> queries per item generated using <code>gemini-2.0-flash</code></li> </ul> Recall at 30 is used as evaluation metric. <code>google/flan-t5-base</code> is used as the generative model to learn to output semantic IDs given the context (be it search or user history) using supervised fine-tuning. <h2 id="results-8"><a class="header" href="#results-8">Results</a></h2> The results show that: <ul> <li>Single-task based embeddings perform best in their own task, much better than any multi-task method. But they degrade performance on the other task very badly.</li> <li>Amongst the multi-task methods, the multi-task method of training a bi-encoder on both contexts works best.</li> <li>Amongst semantic ID tokenisation methods, <code>RQ-Kmeans</code> was the clear winner, far outperforming all other methods. <code>RQ-VAE</code> in particular showed degenerate results especially for the search case.</li> </ul> <blockquote> Question: Why did RQVAE perform so badly compared to the naive RQ-Kmeans? Did the authors handle the RQVAE correctly? </blockquote> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="yan-2025---llm-for-recsys"><a class="header" href="#yan-2025---llm-for-recsys">Yan 2025 - LLM for Recsys</a></h1> <a href="https://www.youtube.com/watch?v=2vlCqD6igVA">Recsys Keynote: Improving Recommendation Systems & Search in the Age of LLMs - Eugene Yan, Amazon</a> This talk covers future of recsys and how LLMs can be incorporated. 3 challenges: <ul> <li>Cold start challenge of hash based item IDs</li> <li>Lack of metadata</li> <li>Task specific models duplicate engineering, increase maintenance cost and don't benefit from transfer learning <ul> <li>Benefits: simplifies systems, reduces maintenance and transfer learning</li> <li>But there may be alignment tax</li> </ul> </li> </ul> <h2 id="kuaishou-example-for-semantic-ids"><a class="header" href="#kuaishou-example-for-semantic-ids">Kuaishou Example for semantic IDs</a></h2> <ul> <li>Challenge: hash based item IDs don't encode item content and struggle with cold start and sparsity problem.</li> <li>Solution: Semantic IDs based on multimodal content</li> </ul> Kuaishou is a short video platform. The main problem they wanted to tackle is to help users discover new items faster. Idea: <ul> <li>Train standard ID-based embeddings for user and items</li> <li>Create cluster ID from concatenated content embeddings <ul> <li>Text: BERT</li> <li>Video: ResNet</li> <li>Audio: VGGish</li> </ul> </li> <li>Run k-means on <code>100 million</code> items into around <code>1k</code> clusters <ul> <li>Each cluster gets an ID and also an embedding</li> <li>Incorporate cluster embedding in final embedding</li> </ul> </li> </ul> Result: <ul> <li> <code>+3.4% clicks</code> </li> <li> <code>+3.0% likes</code> </li> <li> <code>+3.6% cold start coverage</code> (% of item impressions which are new items) </li> <li> <code>+1.2% cold start velocity</code> (% of new items that were able to hit some threshold of views) </li> <li> Example: <ul> <li>trainable, multimodal, semantic IDs @ Kuaishou</li> <li>Short videos platform</li> <li>Problem: help users discover new items faster</li> </ul> </li> </ul> <h2 id="filtering-bad-job-recommendations-at-indeed"><a class="header" href="#filtering-bad-job-recommendations-at-indeed">Filtering Bad job recommendations at Indeed</a></h2> Problem: poor user experience of email job recommendations and lost trust due to low quality job recommendations Solution: Lightweight classifier trained from GPT-4o annotated data to filter bad recs Process: <ul> <li>Start with evals - 250 labelled examples with confidence labels</li> <li>Started with open LLMs like Llama2, but performance was very bad</li> <li>Used GPT-4, which was very accurate but too slow and costly (22 secs)</li> <li>Used GPT-3.5, but had poor precision (0.63) on job recommendations</li> <li>Finetuned GPT-3.5 and got 0.9 precision at 1/4th of of GPT-4 cost and latency</li> <li>Distilled lightweight classifier on finetuned GPT-3.5 labels</li> </ul> Lightweight classifier was <code>0.86 auc-roc</code>, with latency <code><200ms</code>. Result was: <ul> <li><code>-18% bad recommendations</code> <ul> <li>Expected lower application rates because recommending fewer items</li> </ul> </li> <li><code>unsubscribed rate -5%</code></li> <li><code>application rate +5%</code></li> </ul> <h2 id="enriching-exploratory-search-queries--spotify"><a class="header" href="#enriching-exploratory-search-queries--spotify">Enriching exploratory search queries @ Spotify</a></h2> Problem: Help users search for new items (podcasts, audiobooks) in a catalogue of known items (e.g. songs, artists) <ul> <li>How to solve cold start issue for new categories?</li> <li>Exploratory search was essential to expand beyond music</li> </ul> Solution: Query recommendation system Start creating queries from new items (e.g. podcast title, author etc.) and ask LLM to rewrite as natural language query <h2 id="unified-ranker-for-search--recsys--netflix"><a class="header" href="#unified-ranker-for-search--recsys--netflix">Unified Ranker for Search & Recsys @ Netflix</a></h2> <a href="https://www.arxiv.org/abs/2408.10394">Joint Modeling of Search and Recommendations Via an Unified Contextual Recommender (UniCoRn)</a> Example of Stripe building a transformer based foundation model from sequence of transactions to identify fraud. Problem: teams deal with complexity from bespoke models for search, similar item recs, pre-query recs <ul> <li>High operational cost and missed transfer learning opportunities</li> </ul> Unified Contextual Ranker (UniCoRn) takes in a unified input schema and returns a prediction. Unified inputs: <ul> <li>User ID</li> <li>item ID</li> <li>Search Query</li> <li>Task</li> </ul> Some clever tricks to reframe item to item recommendations as search, by using last item title as query. Unified model used for search, pre-query filtering, video to video recs and more. Able to match or exceed previous task based models. Can iterate much faster with a unified model. <h2 id="unified-embeddings--etsy"><a class="header" href="#unified-embeddings--etsy">Unified Embeddings @ Etsy</a></h2> Problem: How to help users get better results with highly specific or broad queries, on ever-changing inventory. <ul> <li>Query <code>mother's day gift</code> does not match product vocabulary</li> <li>Lexical retrieval does not account for user preferences</li> </ul> Solution: Unified embedding and retrieval model <ul> <li>Two tower architecture for user and product side</li> <li>Add a quality vector on the product side (rating, freshness, conversion rate) concatenated to the product vector</li> <li>Add a constant vector on the user side just to make dimensions match</li> </ul> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="tandon-2025---gemini-for-youtube"><a class="header" href="#tandon-2025---gemini-for-youtube">Tandon 2025 - Gemini for YouTube</a></h1> <a href="https://www.youtube.com/watch?v=LxQsQ3vZDqo">Teaching Gemini to Speak YouTube: Adapting LLMs for Video Recommendations to 2B+DAU - Devansh Tandon</a> In general, recommendations is going to have higher reach than search for most consumer apps. The personalized recommendation problem is to learn <code>f(user, context) = recs</code>. How to rethink the whole recommendation process in terms of gemini? <ul> <li>Large Recommender Model (LRM): adapting gemini for recommendation tasks</li> <li>Start with a base gemini checkpoint, adapted for youtube recommendations</li> <li>Align LRM to different tasks using adapters <ul> <li>Video retrieval (home, watch next, shorts, search, music)</li> <li>Video ranking</li> </ul> </li> </ul> To train LRM, need to develop video tokens <ul> <li>Each video is one token, so we can represent watch history</li> <li>And then output some sequence of video recs</li> </ul> SemanticID: tokenize videos to create a language for Youtube <ul> <li>Creating the atomic language for youtube videos, move away from hash based tokenizing</li> <li><code>Youtube video -> extract features like title, description, transcript, audio, vidoe frames -> embedding -> RQ-VAE to quantize embedding into semantic ID</code></li> <li>E.g. a semantic ID would look like <code>A228 B204 C196 D413 E589 ...</code></li> </ul> Continued pre-training: get LRM to understand both English and the semantic ID "language": <ul> <li>Synthetic data relationships: <ul> <li>Link text and SID, e.g. <code>Prompt: Video [A228 ...] has title</code> and the response is <code>Carlos Alcaraz vs Novak Djokovic</code></li> <li>Create lots of synthetic data doing this</li> </ul> </li> <li>User behaviours: SID sequences of video watches <ul> <li><code>Prompt: A user has previously watched the following vidoes: [A110, ...], [F707, ...], [C94, ...], \n <mask1></code></li> <li><code>Output: <mask1>: [B230, ...]</code></li> </ul> </li> </ul> After doing this kind of training, we get an LRM that can reason across both English and Youtube Video language. <ul> <li>e.g. <code>Prompt: Video [A185, ...] is interesting to tennis fans since it is about wimbledon. Video [J110 ...] is interesting to</code></li> <li>Model can respond <code>Output: scientists since it is about AI</code></li> </ul> We can now perform generative retrieval with this new language. For example, we can just load context and history into the prompt, and get it to recommend new videos. <pre><code>Prompt: ------ User: region US | 24 years female | device ANDROID | origin watch next| Context video: channel Olympics | title WHAT a COMEBACK! Men's 400m | SemanticID_1 | Watch history: SID 1 Taylor Swift 100% 260.00s SID 2 Kris Hui 40% 260.00s </code></pre> Interestingly, compared to traditional recommenders which returns mostly men's sports videos (due to the last video being <code>Men's 400m</code>), the generative system was able to recommend <code>women's sprint events</code>. In general: <ul> <li>The LRM is able to learn very quickly, very data efficient</li> <li>Handles the toughest recs when user information is very limited</li> <li>But servicing costs are very large for YouTube scale, so alot of effort focused on reducing TPU serving cost</li> </ul> A simple trick to reduce serving cost is to do batch recommendations <ul> <li>Build a simple video to recs table where given a seed video, get the LRM to output generic recommendations</li> <li>Prompt is something like <code>language {seed_lang} | duration {video_length} | age {video_age} | title {title} | channel {uploader} | {seed_sid}</code></li> <li>Everyday, take <code>top 20M</code> of videos by views in <code>last 7days</code>, and do batch inference</li> <li>Use these batch recommendations to serve some users</li> </ul> Base Gemini is too large at YouTube scale, have to use smaller more efficient models to adapt for recsys. <blockquote> LLM x RecSys Recipe <ol> <li>Tokenize content</li> </ol> <ul> <li>Capture the content essence into an atomic, semantic token</li> </ul> <ol start="2"> <li>Adapt LLM: english <> domain language</li> </ol> <ul> <li>Adapt foundation model to reason across english & new tokens (bilingual)</li> </ul> <ol start="3"> <li>Prompt with user information</li> </ol> <ul> <li>Construct prompts with user information, activity, actions</li> <li>Train surface/task-specific models</li> </ul> </blockquote> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="hameed-2025---360brew"><a class="header" href="#hameed-2025---360brew">Hameed 2025 - 360Brew</a></h1> <a href="https://www.youtube.com/watch?v=U0S6CfzAY5c">360Brew: LLM-based Personalized Ranking and Recommendation - Hamed and Maziar, LinkedIn AI</a> Pain points on current LinkedIn ML: <ul> <li>Operational: costly low agility development lifecycle</li> <li>Quality: disjoint optimization</li> <li>Developer experience: slow to roll out changes to models one by one</li> </ul> Goal: build a foundational model capturing the lifetime member activity data that solves all LinkedIn matching problems <ul> <li>Zero shot capability: works well out of the box for next prediction tasks <ul> <li>Measure how well the model does on new products</li> </ul> </li> <li>In-Context learning: Learning from few examples without needing to retrain <ul> <li>How well does the model do on new users / items? [cold start]</li> </ul> </li> <li>Follow instruction from developers / users <ul> <li>User control via prompts</li> </ul> </li> </ul> <h2 id="development"><a class="header" href="#development">Development</a></h2> Building the LLM: <ul> <li>Need to convert user history into a prompt by verbalizing user information and activities</li> <li>Provide instruction on what problem we are solving</li> <li>At time of training use different styles for verbalization</li> </ul> Prompt looks something like: <pre><code class="language-markdown" style="white-space: pre-wrap;"> ## Instruction You are provided a member's profile and a set of jobs, their description, and interactions that the member had with the jobs. For each past job, the member has taken one of the following actions: applied, viewed, dismissed, or did not interact. Your task is to analyze the job interaction data along with the member's profile to predict whether the member will apply, view, or dismiss a new job referred to as the "Question" job. Note: Focus on skills, location, and years of experience more than other criteria. ## Member Profile Current position: software engineer, current company: LinkedIn, Location: Sunnyvale, California. ## Past job interaction data Member has applied to the following jobs: [Age: 2 days, Title: Software Engineer, Location: New York, Country: USA, Company: Meta, Description: . . . ] Member has viewed the following jobs: [Age: 1 week, Title: Software Engineer, Location: Texas, Country: USA, Company: AMD, Description: . . . ] ## Question 1 Will the member apply to the following job: [Age: 1 day, Title: Software Engineer, Location: Seattle, Country: USA, Company: Apple, Description: . . . ] ## Question 2 Will the member apply to the following job: [Age: 5 days, Title: RF Engineer, Location: Bay Area, Country: USA, Company: Google, Description: . . . ] </code></pre> So in contrast to YouTube's semantic IDs, LinkedIn encodes past interactions in textual form. Development pipeline: <ul> <li>Start with OSS model</li> <li>Continued pre-training</li> <li>Supervised Finetuning</li> <li>Alignment</li> <li>Generate Brew-XL <code>150B</code> model</li> <li>Distill to <code>Brew-mini</code></li> <li>Prune and quantize to <code>Brew-mini-turbo</code> at <code>3B</code> parameters <ul> <li>Ablation studies show that it is critical to <code>first go BIG, then go small</code></li> </ul> </li> </ul> To make development cycle smooth, build in a lot of automation into the pipelines. Especially evaluation loop. Three levers to improve model quality: <ul> <li>More (and better data) <ul> <li>Prepare data to maximize accuracy, distribution of different type of data</li> </ul> </li> <li>Bigger model size</li> <li>Context length <ul> <li>Longer context length means deeper user activity</li> <li>Increasing context length initially improves performance to a certain point (around <code>20-30k tokens</code>)</li> <li>Beyond that models don't generalize that well and performance degrades</li> </ul> </li> </ul> <h2 id="tasting"><a class="header" href="#tasting">Tasting</a></h2> Performance of model is best for cold start users. Measure relative gain over production model: <ul> <li><code>5 maximum activities</code>: <code>+6%</code></li> <li><code>10 maximum activities</code>: <code>+4%</code></li> <li><code>100 maximum activities</code>: <code>+2%</code></li> </ul> Generalization to new domain. 360Brew model can generalize to out of domain tasks and surfaces and beat production models in those tasks. <ul> <li>Increases team agility to roll out new features without training new model</li> </ul> <h2 id="serving"><a class="header" href="#serving">Serving</a></h2> Three levers to improve efficiency: <ul> <li>Sparsification</li> <li>Smaller model <ul> <li>Distillation from big model to small model done using SFT + KD loss</li> <li>Gradual distillation is more effective than direct distillation, i.e. go 8B model to 6B model to 3B model etc.</li> <li>Pruning is done layerwise, gradual pruning</li> </ul> </li> <li>Quantization: Mix precision <ul> <li>FP8 for all weights</li> <li>FP32 for language model head and logit processor. Important for recommendations, otherwise predictions collapse.</li> </ul> </li> <li>Sparsification <ul> <li>Star attention (reduce attention quadratic cost) <ul> <li>Not every item needs to attend to every item</li> </ul> </li> <li>When scoring, we can score multiple items at the same time (sounds like <code>500</code>) <ul> <li>Need to make sure these items do not attend to each other</li> </ul> </li> </ul> </li> </ul> Q&A: <ul> <li>They use 50-60 tasks out of domain to measure the effectiveness of the model in the eval loop.</li> <li>Designed custom vLLM kernels to allow multi-item scoring by modifying the attention mask</li> </ul> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="han-2025---grpo-example"><a class="header" href="#han-2025---grpo-example">Han 2025 - GRPO Example</a></h1> <a href="https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/gpt-oss-(20B)-GRPO.ipynb#scrollTo=MXPZ1MsUMqNZ">Notebook Tuning GPT-OSS</a> This notebook shows how to finetune <code>gpt-oss-20b</code> using unsloth. The parameters used are: <pre><code class="language-python">max_seq_length = 768 lora_rank = 4 lora_alpha = lora_rank * 2 load_in_4bit = True offload_embedding = True </code></pre> The toy problem is to induce gpt-oss to learn how to output an optimized all python <code>matmul</code> procedure that only uses the standard library. Naturally, using optimized libraries like <code>numpy</code> will be a lot faster, but there is some room to improve from a naive implementation in pure python. First, we generate some random matrices. Note that we need <code>A_list</code> and <code>B_list</code> which are <code>list[list[float]]</code> to pass to our pure python function. <pre><code class="language-python">import numpy as np def generate_random_matrices(seed = 3407, n = 256): random_state = np.random.RandomState(seed) n, k, m = random_state.randint(1, n+1, size = 3) A = np.random.uniform(-10, 10, size = (n, k)) B = np.random.uniform(-10, 10, size = (k, m)) return A, A.tolist(), B, B.tolist() </code></pre> For example, a kernel generated by GPT-5 is: <pre><code class="language-python"># Kernel generated by GPT-5 def matmul(A, B): B_transpose = list(zip(*B)) return [ [ sum(a*b for a, b in zip(row, col)) for col in B_transpose ] for row in A ] </code></pre> <h2 id="preventing-cheating"><a class="header" href="#preventing-cheating">Preventing Cheating</a></h2> GRPO works using RL with a reward function. To get the desired behaviour, we have to find a reward function that disincentivises cheating. For example, the LLM might just come up with a solution that imports <code>torch</code> or <code>numpy</code> to get an optimized matmul function. Firstly, we write a function called <code>check_only_stdlib_imports</code> that, given a string representing a python function (let's call it <code>fn_string</code>), checks if it imports anything outside of the python standard library. We omit the definition here as it is rather involved. Secondly, given <code>fn_string</code>, we want to compile it into a python function and disallow the function from importing anything from the global name space. The following code does so: <pre><code class="language-python">def create_locked_down_function(function: str): output_function = {} exec(function, {}, output_function) new_matmul = output_function["matmul"] new_matmul = types.FunctionType(new_matmul.__code__, {}) return new_matmul </code></pre> Some explanations: <ul> <li><code>exec(string: str, globals: dict, locals: dict)</code> <ul> <li>This executes code in <code>string</code> and stores whatever definitions within the code into <code>locals</code></li> <li>By specifying an empty <code>dict</code> for <code>globals</code>, we are disallowing the functions defined within to access the global namespace</li> </ul> </li> <li><code>types.FunctionType(code, globals)</code> <ul> <li>This creates a new function from compiled python bytecode</li> <li>It also redefines the global variables that the function can access as an empty dictionary</li> </ul> </li> <li>So this function turns <code>fn_string</code> into a python function, but disallows accessing any global variables and libraries</li> </ul> <h2 id="benchmarking"><a class="header" href="#benchmarking">Benchmarking</a></h2> Now we create the actual benchmarking function. There are some details here that we shall gloss over, but note that at a high level, we call <code>Benchmarker.benchmark(mat_mul_fn, [A, B])</code> to get some statistics about the matrix multiplcation function we are testing. <pre><code class="language-python">import os, gc, time, statistics import signal from contextlib import contextmanager class TimeoutError(Exception): pass @contextmanager def time_limit(seconds): def _handler(signum, frame): raise TimeoutError(f"Timed out after {seconds}s") old = signal.signal(signal.SIGALRM, _handler) signal.setitimer(signal.ITIMER_REAL, seconds) try: yield finally: signal.setitimer(signal.ITIMER_REAL, 0.0) signal.signal(signal.SIGALRM, old) class Benchmarker: def __init__(self, trials = 3, loops = 1, timeout = 30): self.buffer = np.zeros(2 * 1024 * 1024 * 1024, dtype = np.uint8) self.trials = trials self.loops = loops assert timeout > 0 # Cannot be 0 since it won't work! self.timeout = timeout def thrash(self): # Edit the buffer to wipe cache lines self.buffer ^= 1 return int(self.buffer[::4096].sum()) def benchmark(self, function, arguments): assert len(arguments) == self.loops samples = [] exceptions = [] timed_out = 0 for _ in range(self.trials): gc.collect(); gc.disable(); self.thrash() t_start = time.perf_counter_ns() for i in range(self.loops): try: with time_limit(self.timeout): function(*arguments[i]) except TimeoutError as e: timed_out += 1 except Exception as e: exceptions.append(str(e)) t_end = time.perf_counter_ns() gc.enable() samples.append((t_end - t_start) // max(1, self.loops)) return { "median_ns": int(statistics.median(samples)), "mean_ns": int(statistics.fmean(samples)), "stdev_ns": int(statistics.pstdev(samples) if len(samples) > 1 else 0), "exceptions" : exceptions, "timeouts" : timed_out, } </code></pre> Some other details: <ul> <li>A timeout is added to raise an error if the function runs for too long</li> <li><code>gc</code> is handled to avoid garbage collection from messing up the trial timing</li> <li>Some thrashing of a large <code>2GB</code> buffer is done to wipe out the L1, L2, L3 CPU caches to avoid the CPU from tapping on the cache to fast track the computation</li> </ul> <h2 id="generate-matmul-function"><a class="header" href="#generate-matmul-function">Generate matmul function</a></h2> Now we get gpt-oss to generate some matmul function, like so. <pre><code class="language-python">from transformers import TextStreamer prompt = """ Create a new fast matrix multiplication function using only native Python code. You are given a list of list of numbers. Output your new function in backticks using the format below: ```python def matmul(A, B): return ... ``` """.strip() text = tokenizer.apply_chat_template( [{"role": "user", "content": prompt}], tokenize = False, add_generation_prompt = True, reasoning_effort = "low", ) _ = model.generate( **tokenizer(text, return_tensors = "pt").to("cuda"), temperature = 1.0, max_new_tokens = 512, streamer = TextStreamer(tokenizer, skip_prompt = False), ) </code></pre> Which will give something like: <pre><code class="language-bash">... ```python def matmul(A, B): """ Multiply two matrices A and B where both are given as lists of lists (rows) and are compatible for multiplication. Returns the resulting matrix as a list of lists. """ # Transpose B so we can access its columns as rows. B_T = [list(col) for col in zip(*B)] # O(n*m) time result = [] for row in A: # for each row in A res_row = [] for col in B_T: # for each column of B # Compute dot product of row and column dot = sum(a * b for a, b in zip(row, col)) res_row.append(dot) result.append(res_row) return result ``` ... </code></pre> <h2 id="reward-functions"><a class="header" href="#reward-functions">Reward functions</a></h2> Now we want to build up our reward function which assigns a reward score to each new generation from gpt-oss. The flow is: <ul> <li>Extract the function from a generation text</li> <li>Use <code>create_locked_down_function</code> to turn it into python code</li> <li>Use <code>Benchmarker</code> to test how fast it is</li> <li>Using benchmark statistics, assign some scores</li> </ul> First we define <code>extract_function(generation_text: str)</code> which simply extracts the function from a generation text by looking for the triple-backticks. Omitted for brevity. Next we create multiple reward functions that test different aspects. For example, we have <code>function_works</code> which just checks that the function runs without errors. Note that the function runs on a list of completions. <pre><code class="language-python">def function_works(completions, **kwargs): scores = [] for completion in completions: score = 0 response = completion[0]["content"] function = extract_function(response) print(function) if function is not None: ok, info = check_only_stdlib_imports(function) if function is None or "error" in info: score = -2.0 else: try: new_matmul = create_locked_down_function(function) score = 1.0 except: score = -0.5 scores.append(score) return scores </code></pre> Another reward function checks that the function is correct. As we can see, the score is higher for more accurate functions and lower for inaccurate functions. <pre><code class="language-python">def correctness_check(completions, **kwargs): scores = [] # Generate some random matrices of size less than 128 A, A_list, B, B_list = generate_random_matrices(seed = np.random.randint(10000), n = 128) for completion in completions: score = 0 response = completion[0]["content"] function = extract_function(response) if function is not None: ok, info = check_only_stdlib_imports(function) if function is None or "error" in info: scores.append(0) continue try: new_matmul = create_locked_down_function(function) except: scores.append(0) continue try: pred = new_matmul(A_list.copy(), B_list.copy()) except: # Failed! scores.append(-2.0) continue true = np.matmul(A, B) amax_error, mse_error = calculate_difference(pred, true) # Check correctness and score! machine_epsilon = 100*np.finfo(np.float64).eps if amax_error >= 3: score = -3.0 elif amax_error >= 2: score = -2.5 elif amax_error >= 1: score = -2.0 elif amax_error >= 0.5: score = -1.0 elif amax_error >= 100*machine_epsilon: score = 0.0 elif amax_error >= machine_epsilon: score = 1.0 else: score = 3.0 if mse_error >= 3: score += -3.0 elif mse_error >= 2: score += -2.5 elif mse_error >= 1: score += -2.0 elif mse_error >= 0.5: score += -1.0 elif mse_error >= 100*machine_epsilon: score += 0.0 elif mse_error >= machine_epsilon: score += 1.0 else: score += 3.0 scores.append(score) return scores </code></pre> Finally we have the speed check function which rewards fast implementations. <ul> <li>If the LLM implementation is faster than numpy implementation, then <code>positive</code> is used (e.g. 2x faster gives <code>+0.02</code>)</li> <li>If the LLM implementation is slower than numply implementation, then <code>negative</code> is used (e.g. 2x slower gives <code>-0.02</code>)</li> </ul> <pre><code class="language-python">import gc def speed_check(completions, **kwargs): scores = [] # Generate some random matrices of size less than 256 A, A_list, B, B_list = generate_random_matrices(seed = np.random.randint(10000), n = 256) numpy_results = benchmarker.benchmark(np.matmul, [(A, B)]) for completion in completions: score = 0 response = completion[0]["content"] function = extract_function(response) if function is not None: ok, info = check_only_stdlib_imports(function) if function is None or "error" in info: scores.append(0) continue try: new_matmul = create_locked_down_function(function) except: scores.append(0) continue new_results = benchmarker.benchmark(new_matmul, [(A_list.copy(), B_list.copy())]) # Get score and clip to -10, 10 negative = -(new_results["median_ns"] / numpy_results["median_ns"]) / 100 positive = +(numpy_results["median_ns"] / new_results["median_ns"]) / 100 score = negative if new_results["median_ns"] >= numpy_results["median_ns"] else positive if score >= 10: score = 10 if score <= -10: score = -10 scores.append(score) # Free memory to counteract OOMs gc.collect() torch.cuda.empty_cache() return scores </code></pre> <h2 id="grpotrainer"><a class="header" href="#grpotrainer">GRPOTrainer</a></h2> Finally, we can prepare our dataset and set up the <code>GRPOTrainer</code>. First, we create a dataset with just the <code>prompt</code>. We just need the prompt because we're running this in dynamic mode, where <code>GRPOTrainer</code> will sample a few generations for each prompt and evaluate it using the reward functions. This represents true on-policy RL. Note also that we use <code>reasoning_effort="low"</code>. <pre><code class="language-python">from datasets import Dataset dataset = Dataset.from_list([{"prompt" : [{"role": "user", "content": prompt.strip()}], "answer" : 0, "reasoning_effort": "low"}]*1000) maximum_length = len(tokenizer(prompt.strip())["input_ids"]) print(maximum_length) dataset[0] </code></pre> Now we setup <code>GRPOTrainer</code> with configuration. <pre><code class="language-python">max_prompt_length = maximum_length + 1 # + 1 just in case! max_completion_length = max_seq_length - max_prompt_length from trl import GRPOConfig, GRPOTrainer training_args = GRPOConfig( temperature = 1.0, learning_rate = 5e-5, weight_decay = 0.01, warmup_ratio = 0.1, lr_scheduler_type = "linear", optim = "adamw_8bit", logging_steps = 1, per_device_train_batch_size = 1, gradient_accumulation_steps = 1, # Increase to 4 for smoother training num_generations = 2, # Decrease if out of memory max_prompt_length = max_prompt_length, max_completion_length = max_completion_length, # num_train_epochs = 1, # Set to 1 for a full training run max_steps = 100, save_steps = 100, report_to = "none", # Can use Weights & Biases output_dir = "outputs", # For optional training + evaluation # fp16_full_eval = True, # per_device_eval_batch_size = 4, # eval_accumulation_steps = 1, # eval_strategy = "steps", # eval_steps = 1, ) </code></pre> Finally we can train the model. Note that T4 GPU might take 5 minutes for one generation which might take hours to start seeing some progress on this task. <pre><code class="language-python"># For optional training + evaluation # new_dataset = dataset.train_test_split(test_size = 0.01) trainer = GRPOTrainer( model = model, processing_class = tokenizer, reward_funcs = [ function_works, no_cheating, correctness_check, speed_check, ], args = training_args, train_dataset = dataset, # For optional training + evaluation # train_dataset = new_dataset["train"], # eval_dataset = new_dataset["test"], ) trainer.train() </code></pre> After training, we can use our original prompt to see the new function that the LLM generates. <pre><code class="language-python">text = tokenizer.apply_chat_template( [{"role": "user", "content": prompt}], tokenize = False, add_generation_prompt = True, reasoning_effort = "low", ) from transformers import TextStreamer _ = model.generate( **tokenizer(text, return_tensors = "pt").to("cuda"), temperature = 1.0, max_new_tokens = 1024, streamer = TextStreamer(tokenizer, skip_prompt = False), ) </code></pre> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="yan-2025---semantic-ids"><a class="header" href="#yan-2025---semantic-ids">Yan 2025 - Semantic IDs</a></h1> <ul> <li>Article: <a href="https://eugeneyan.com/writing/semantic-ids/">https://eugeneyan.com/writing/semantic-ids/</a></li> <li>Code: <a href="https://github.com/eugeneyan/semantic-ids-llm">https://github.com/eugeneyan/semantic-ids-llm</a></li> </ul> Eugene Yan's experience of training an LLM-recommender hybrid that can recommend items using natural language. <h2 id="extend-the-vocab"><a class="header" href="#extend-the-vocab">Extend the vocab</a></h2> Insert semantic ID tokens like <code><|sid_0|></code>, <code><|sid_1|></code> into the vocab. Sequences of these tokens will be used to represent the catalog. So an item may be represented like: <pre><code><|sid_start|><|sid_173|><|sid_324|><|sid_764|><|sid_768|><|sid_end|> </code></pre> <h2 id="data"><a class="header" href="#data">Data</a></h2> He used video games category of Amazon Reviews 2023, because it has rich product metadata and lots of behavioural data. <ul> <li>Keep only products with titles longer that <code>20 chars</code> and description longer than <code>100 chars</code></li> <li><code>66k</code> products</li> <li><code>737k</code> rows of interactions</li> </ul> <h2 id="training-semantic-ids"><a class="header" href="#training-semantic-ids">Training Semantic IDs</a></h2> Used RQ-VAE to train semantic IDs. Will not expound on RQ-VAE here. Eugene used: <ul> <li><code>3-level</code> codebook</li> <li>Each level has <code>256</code> codes</li> <li>This has collisions on around <code>10%</code> of the <code>66k products</code></li> <li>Added a sequentially increasing token to each ID to ensure uniqueness</li> </ul> <h2 id="baselines"><a class="header" href="#baselines">Baselines</a></h2> <ol> <li>Train a SASRec model on semantic IDs to compare against traditional SASRec</li> <li>Use Qwen3-Embedding-0.6B to encode product metadata into embeddings</li> <li>Finetune Qwen3-8B to recommend items via semantic IDs</li> </ol> <h2 id="data-cleaning"><a class="header" href="#data-cleaning">Data Cleaning</a></h2> <ul> <li>Clean item descriptions using Gemini 2.5 Flash to remove HTML, reduce verbosity</li> <li>Remove promotional text, standardize formatting for titles etc.</li> <li>Augment data by extracting structured metadata like product type, platform, genre, hardware type etc.</li> <li>Build user interaction histories</li> </ul> <h2 id="training-rq-vae"><a class="header" href="#training-rq-vae">Training RQ-VAE</a></h2> First, we get the <code>1,024</code> embeddings using Qwen3-0.6B, then l2-normalize. The RQ VAE consists of: <ul> <li>An encoder</li> <li>A codebook with 3 quantization levels</li> <li>A symmetric decoder</li> </ul> Some tricks: <ul> <li>Use rotation trick as a replacement for the straight through estimator</li> <li>Initialize codebooks with k means clustering</li> <li>Reset unused codes</li> <li>Large batch size</li> </ul> Metrics for measuring quality of RQ-VAE: <ol> <li><code>loss/reconstruction</code>: How well we can reconstruct the original item embeddings after compressing and decompressing</li> <li><code>loss/vq</code>: Combined codebook and commitment loss across all levels. <ul> <li>Ensures that encoder outputs and codebook vectors stay close together</li> </ul> </li> <li><code>loss/total</code>: sum of <code>loss/vq</code> and <code>loss/reconstruction</code></li> <li><code>loss/validation</code>: total loss but on the held out validation set. Our deciding metric.</li> <li><code>metrics/avg_residual_norm</code>: Leftover residuals between the quantized embedding and the original embedding</li> <li><code>metrics/unique_ids_proportion</code>: % of items with unique IDS in a batch <ul> <li>Helps to check against codebook collapse</li> <li>We want this metric to be high</li> </ul> </li> <li>Codebook distribution <ul> <li>Plot the item distribution amongst the <code>256</code> codes at each level</li> <li>It should look like a uniform distribution</li> </ul> </li> </ol> Some hyperparameter tuning: <ul> <li>Set β=0.5: <ul> <li>Tried β=0.25,0.5,1.0. β is the commitment weight that balances reconstruction accuracy with codebook commitment</li> <li>β=0.5 had the lowest validation loss, so it was chosen</li> </ul> </li> <li>Investing in data cleaning significantly improved all the metrics</li> </ul> <h2 id="sasrec-comparison"><a class="header" href="#sasrec-comparison">SASRec comparison</a></h2> Eugene tested two SASRec variants: <ul> <li>Traditional SASRec <ul> <li>Each item has a unique ID with no semantic meaning</li> <li>Model uses 2 causal attention blocks, 64 dim hidden dimension, trained with binary cross entropy loss</li> <li>Dot product of embeddings is used to generate scores</li> </ul> </li> <li>Semantic SASRec <ul> <li>For semantic version, each item is a sequence of semantic tokens</li> <li>Hence for each position, the model needs to generate a sequence of tokens to represent an item</li> <li>Instead of binary cross entropy loss, we need to sum up the cross entropy loss at each position</li> <li>Question: does it make sense to add some weights to put more weightage on the first semantic token and decay it for subsequent positions?</li> <li>Teacher forcing is used for this training</li> <li>A larger model is used here, 4 causal attention blocks and 384 hidden dim</li> </ul> </li> </ul> From this experiment, Eugene found that traditional SASRec is significantly better. But he puts this down to the difficulty of generating a sequence of tokens compared to directly generating one token. I also note that we are not using a pretrained LLM here, which means we are missing out on some pretrained capabilities that we could have tapped on. <h2 id="finetuning-qwen-8b"><a class="header" href="#finetuning-qwen-8b">Finetuning Qwen-8B</a></h2> Now finally we train the language model to converse in semantic IDs. First, we build the training dataset of <code>4.2 million</code> conversations of various task types: <ul> <li>Given a semantic ID, predict the item description</li> <li>Given item description, predict the semantic ID</li> <li>Predict the next item in the user's sequence</li> <li>Understand relationships between item categories</li> <li>Multi-hop reasoning</li> </ul> Each of these are formatted as a conversation with a system prompt, user prompt and assistant response. Now, we finetune the model in two phases. In Phase 1, we focus on extending the model's vocabulary. <ul> <li>We add <code>1,027</code> new vocabulary items to the model's vocabulary, and resize the model's embedding matrix for them.</li> <li>In this phase, we freeze all the model parameters except the input and output embedding layers, training <code>15%</code> of total parameters</li> <li>Train for <code>1,000</code> steps with relatively high learning rate</li> </ul> In Phase 2, we do the full training: <ul> <li>Train for 3 epochs on the full dataset</li> <li>3% warmup, lower learning rate, 8-bit AdamW optimizer for memory efficiency</li> <li>Monitor progress with callbacks: <ul> <li>Track gradient norms</li> <li>Measure hierarchical accuracy (I guess on different codebook granularities) on recommendation tasks</li> <li>Sample model outputs on fixed set of conversational inputs</li> </ul> </li> </ul> <h2 id="inference"><a class="header" href="#inference">Inference</a></h2> Now the model is trained, we can have chat / recommendation conversations with it: <blockquote> User: “Halo 3 Limited Edition -Xbox 360”, “Lost Planet: Extreme Condition - Xbox 360”, “Lost Planet: Extreme Condition Collector’s Edition -Xbox 360” <|rec|> Model: <code><|sid_start|><|sid_64|><|sid_351|><|sid_733|><|sid_768|><|sid_end|></code> </blockquote> The model can also explain its recs: <blockquote> User: I just finished “Dragon Quest Heroes II Explorer’s Edition - PlayStation 4”. Suggest another <|rec|> and explain why: Model: <code><|sid_start|><|sid_139|><|sid_289|><|sid_534|><|sid_768|><|sid_end|></code> If you liked Dragon Quest Heroes II, you might like Nights of Azure because both are action RPGs for the PlayStation 4 with a focus on combat and character progression. Both games offer a narrative-driven experience with a strong emphasis on combat mechanics, suggesting a shared player interest in this genre. </blockquote> We can even have multi-turn conversation with the model: <blockquote> User: I’m looking for games similar to Mario Kart. <|rec|> Model: <code><|sid_start|><|sid_131|><|sid_503|><|sid_738|><|sid_768|><|sid_end|>, <|sid_start|><|sid_131|><|sid_260|><|sid_577|><|sid_768|><|sid_end|>, <|sid_start|><|sid_145|><|sid_416|><|sid_671|><|sid_768|><|sid_end|></code> User: How about something similar but for Xbox? <|rec|> Model: <code><|sid_start|><|sid_183|><|sid_461|><|sid_517|><|sid_768|><|sid_end|>, <|sid_start|><|sid_183|><|sid_461|><|sid_562|><|sid_768|><|sid_end|>,<|sid_start|><|sid_183|><|sid_461|><|sid_616|><|sid_768|><|sid_end|></code> User: Suggest a name and description for the bundle. Model: Xbox Racing Legends: Speed & Style Pack </blockquote> <h2 id="code-1"><a class="header" href="#code-1">Code</a></h2> Here we deep dive into Eugene's code and how it is implemented. Most of the code is contained in the <a href="https://github.com/eugeneyan/semantic-ids-llm/tree/main/src"><code>/src</code> directory</a>. <ul> <li><a href="articles/yan_2025_semantic_ids.html#device_managerpy">device_manager.py</a></li> <li><a href="articles/yan_2025_semantic_ids.html#tokenize_itemspy">tokenize_items.py</a></li> <li><a href="articles/yan_2025_semantic_ids.html#embed_itemspy">embed_items.py</a></li> <li><a href="articles/yan_2025_semantic_ids.html#train_rqvaepy">train_rqvae.py</a> <ul> <li><a href="articles/yan_2025_semantic_ids.html#quantizationoutput">QuantizationOutput</a></li> <li><a href="articles/yan_2025_semantic_ids.html#vectorquantizer">VectorQuantizer</a></li> <li><a href="articles/yan_2025_semantic_ids.html#rqvae">RQVAE</a></li> </ul> </li> <li><a href="articles/yan_2025_semantic_ids.html#finetune_qwen3_8b_vocabpy">finetune_qwen3_8b_vocab.py</a> <ul> <li><a href="articles/yan_2025_semantic_ids.html#finetuneconfig">FineTuneConfig</a></li> <li><a href="articles/yan_2025_semantic_ids.html#extend_tokenizer">extend_tokenizer</a></li> <li><a href="articles/yan_2025_semantic_ids.html#prepare_model">prepare_model</a></li> <li><a href="articles/yan_2025_semantic_ids.html#load_sid_dataset">load_sid_dataset</a></li> <li><a href="articles/yan_2025_semantic_ids.html#datainspectioncallback">DataInspectionCallback</a></li> <li><a href="articles/yan_2025_semantic_ids.html#embeddingmonitorcallback">EmbeddingMonitorCallback</a></li> <li><a href="articles/yan_2025_semantic_ids.html#semanticidgenerationcallback">SemanticIDGenerationCallback</a></li> <li><a href="articles/yan_2025_semantic_ids.html#train_embeddings">train_embeddings</a></li> </ul> </li> </ul> <h3 id="device_managerpy"><a class="header" href="#device_managerpy">device_manager.py</a></h3> The <code>DeviceManager</code> detects the device (<code>cpu</code> or <code>cuda</code> or <code>mps</code>) and is instantiated early in the scripts. The interesting part <code>torch.set_float32_matmul_precision("high")</code> performed when device is <code>cuda</code>. It's supposed to speed up <code>float32</code> operations? <h3 id="tokenize_itemspy"><a class="header" href="#tokenize_itemspy">tokenize_items.py</a></h3> This script tokenizes the product descriptions of video games: <ul> <li>Uses <code>Qwen/Qwen3-Embedding-0.6B</code></li> <li>Batch size = <code>32</code>, max length = <code>2048</code></li> <li>Reads the data from <code>data/output/Video_Games_items_updated.parquet</code> into polars</li> <li>Looks for the <code>item_context</code> field (already preprocessed)</li> <li>Uses the following prompt which will be tokenized</li> </ul> <pre><code>Instruct: Given a product description, generate a semantic embedding that captures its key features and characteristics. Query: {original_item_text} </code></pre> <ul> <li>Saves tokenized <code>input_ids</code> and <code>attention_masks</code> and saves them in <code>.npz</code> (compressed numpy) format using <code>np.savez_compressed</code></li> </ul> <h3 id="embed_itemspy"><a class="header" href="#embed_itemspy">embed_items.py</a></h3> Takes the tokenized file and embeds them. <ul> <li>Writes embedded items into a parquet file</li> </ul> <h3 id="train_rqvaepy"><a class="header" href="#train_rqvaepy">train_rqvae.py</a></h3> Compresses an item embedding into hierarchical semantic IDs. The <code>RQVAEConfig</code> defines the following parameters: <ul> <li><code>item_embedding_dim</code>: embedding dim of our embedding model</li> <li><code>encoder_hidden_dims</code>: <code>[512, 256, 128]</code> the size of the VAE encoder</li> <li><code>codebook_embedding_dim</code>: <code>32</code> Dimension of codebook vectors <ul> <li>Qn: this does not need to match qwen embedding dim?</li> </ul> </li> <li><code>codebook_quantization_levels</code>: <code>3</code> levels in the codebook</li> <li><code>codebook_size</code>: <code>256</code> number of codes per level</li> <li><code>commitment_weight</code>: <code>0.25</code> Commitment loss weight (beta)</li> <li><code>use_rotation_trick</code>: <code>True</code> Use rotation trick for better gradient flow</li> <li><code>batch_size</code>: <code>32768</code> training batch size (why so large?)</li> <li><code>gradient_accumulation_steps</code>: <code>1</code></li> <li><code>num_epochs</code>: <code>20000</code></li> <li><code>scheduler_type</code>: <code>cosine_with_warmup</code></li> <li><code>warmup_start_lr</code>: <code>1e-8</code> used for cosine_with_warmup</li> <li><code>warmup_steps</code>: <code>200</code> used for cosine_with_warmup</li> <li><code>max_lr</code>: <code>3e-4</code> maximum learning rate (start of cosine)</li> <li><code>min_lr</code>: <code>1e-6</code> minimum learning rate (end of cosine)</li> <li><code>use_gradient_clipping</code>: <code>True</code></li> <li><code>gradient_clip_norm</code>: <code>1.0</code></li> <li><code>use_kmeans_init</code>: <code>True</code> initializes codebook vectors using k-means</li> <li><code>reset_unused_codes</code>: <code>True</code> reset unused codes periodically to avoid collapse</li> <li><code>steps_per_codebook_reset</code>: <code>2</code> Reset unused codebook codes every N steps</li> <li><code>codebook_usage_threshold</code>: <code>1.0</code> only reset if usage falls below this proportion (0-1)</li> <li><code>val_split</code>: <code>0.05</code></li> <li><code>steps_per_train_log</code>: <code>10</code> log every N steps</li> <li><code>steps_per_val_log</code>: <code>200</code> validate and checkpoint every N steps</li> </ul> <code>EmbeddingDataset</code> is a torch dataset holding the embeddings: <ul> <li>Extracts all embeddings and holds them in a <code>torch.tensor</code> at init <ul> <li>Not worried about OOM?</li> </ul> </li> </ul> <h4 id="quantizationoutput"><a class="header" href="#quantizationoutput">QuantizationOutput</a></h4> <code>QuantizationOutput</code> is used to hold data for one quantized item: <ul> <li>Holds the local loss for one codebook layer</li> <li>Subclasses <code>NamedTuple</code> which is more lightweight than <code>dataclass</code></li> <li><code>quantized_st: Tensor</code> <ul> <li>The quantized vector which is passed onto the next layer</li> <li>Has the "gradient trick" (either straight through or rotation trick) applied</li> <li>Allows backpropagation into the encoder even though we passed through the non-differentiable codebook layer</li> </ul> </li> <li><code>quantized: Tensor</code> <ul> <li>The raw nearest neighbour vectors from the codebook</li> <li>No gradients</li> <li>Observation <code>quantized</code> and <code>quantized_st</code> should be identical in values, just that one has gradients attached</li> </ul> </li> <li><code>indices: Tensor</code> <ul> <li>These are integer indices which represent the semantic IDs</li> </ul> </li> <li><code>loss: Tensor</code> <ul> <li>The combined loss for this specific codebook layer</li> <li><code>loss = codebook_loss + beta * commitment_loss</code></li> </ul> </li> <li><code>codebook_loss: Tensor</code> <ul> <li>Measures how well the codebook vector matches the encoder output</li> </ul> </li> <li><code>commitment_loss: Tensor</code> <ul> <li>Measures how well the encoder output matches the codebook vectors</li> </ul> </li> </ul> <h4 id="vectorquantizer"><a class="header" href="#vectorquantizer">VectorQuantizer</a></h4> The <code>VectorQuantizer</code> implements one layer of the codebook and is the meat of the logic for training RQVAE. It will be stacked together layer multiple times to form the codebook. <ul> <li>At initialization: <ul> <li>Initializes with <code>RQVAEConfig</code> to hold parameters</li> <li>Initialize <code>self.embedding</code> to an embedding of size <code>codebook_size=256, codebook_embedding_dim=32</code> <ul> <li>Uniform initialization <code>self.embedding.weight.data.uniform_(-1 / self.codebook_size, 1 / self.codebook_size)</code></li> </ul> </li> <li>Registers some buffers for tracking codebook usage: <ul> <li><code>self.register_buffer("usage_count", torch.zeros(self.codebook_size))</code></li> <li><code>self.register_buffer("update_count", torch.tensor(0))</code></li> </ul> </li> </ul> </li> <li><code>find_nearest_codes(x)</code>: <ul> <li>Takes an input vector <code>x</code>, compares it to all vectors in the codebook, and returns the nearest one</li> <li>Simply uses <code>torch.cdist</code> to compute distances, then <code>torch.argmin</code> to get the nearest</li> <li>Returns a tuple of torch tensors: <ul> <li>The nearest index (i.e. codeword) to <code>x</code></li> <li>The quantized embedding at the index position</li> </ul> </li> </ul> </li> <li><code>forward(x)</code> -> <code>QuantizationOutput</code>: <ul> <li>Finds the nearest index and quantized embeddings for a batch of <code>x</code> <ul> <li>Call <code>find_nearest_codes</code> to get <code>indices</code> and <code>quantized</code></li> </ul> </li> <li>Applies the gradient estimator to get <code>quantized_st</code> <ul> <li>This will be used for gradient backprop to the encoder later</li> <li><code>apply_gradient_estimator</code> either uses the straight through or rotation method</li> </ul> </li> <li>Compute losses: <ul> <li><code>codebook_loss</code> is the MSE loss between <code>x.detach()</code> and <code>quantized</code> <ul> <li>We want to pull the codebook embeddings toward <code>x</code></li> </ul> </li> <li><code>commitment_loss</code> is the MSE loss between <code>x</code> and <code>quantized.detach()</code> <ul> <li>We want to pull encoder output toward codebook embeddings</li> </ul> </li> <li><code>loss = codebook_loss + beta * commitment_loss</code></li> </ul> </li> <li>Everything is packaged into <code>QuantizationOutput</code> and returned</li> <li><code>self.update_usage</code> is also called: <ul> <li>Updates counts of which indices were the nearest to <code>x</code></li> <li>Updates the number of training steps</li> </ul> </li> </ul> </li> <li>Straight through <ul> <li>The straight through gradient estimator simply returns <code>x + (quantized - x).detach()</code></li> <li>Essentially, the embeddings passed forward is <code>quantized</code></li> <li>But the vector used for gradient backprop is <code>x</code> (hence straight-through back to the encoded <code>x</code>)</li> <li>This is a naive method but works well enough</li> </ul> </li> <li>Rotation <ul> <li>The problem with the straight through estimator is that we use <code>quantized</code> for the forward pass but use <code>x</code> for the backward pass <ul> <li>This can be problematic especially if <code>q</code> and <code>x</code> are far apart</li> </ul> </li> <li>The rotation idea is to apply a rotation to <code>x</code> until it aligns with <code>q</code> <ul> <li>Since the rotation is differentiable, we get better gradients back to <code>x</code></li> </ul> </li> <li>We compute (to check later): <ul> <li>Let u=x/∣∣x∣∣</li> <li>w=∣∣u+q∣∣u+q is the halfway vector between <code>u</code> and <code>q</code></li> <li>xrot=x−2⟨x,w⟩w+2⟨x,u⟩q</li> </ul> </li> </ul> </li> <li><code>reset_unused_codes</code> <ul> <li>Look up <code>self.usage_count</code> to find unused indices (used <code>0</code> times)</li> <li>Take the current batch of encoded data, and randomly select them to become the new codebook vectors</li> <li>This makes it likely for them to be used in the next forward pass since they correspond to actual encoder outputs</li> <li>All usage counters are reset after this</li> </ul> </li> </ul> <h4 id="rqvae-1"><a class="header" href="#rqvae-1">RQVAE</a></h4> The RQVAE class now assembles multiple <code>VectorQuantizer</code> into the actual VAE to create semantic IDs. At initialization: <ul> <li><code>self.encoder</code>: a simple MLP that shrinks the input embedding down to the codebook dimension <ul> <li>In this code, we go from <code>1024 -> 512 -> 256 -> 128 -> 32</code></li> </ul> </li> <li><code>self.decoder</code>: a simple MLP that goes backward from quantized vector up to embedding dimension <ul> <li>In this code, the decoder dims are just the reversed of the encoder dims</li> <li>So we go from <code>32 -> 128 -> 256 -> 512 -> 1024</code></li> </ul> </li> <li>Both encoder and decoder are wrapped in <code>nn.Sequential</code></li> <li><code>self.vq_layers</code> contains the <code>VectorQuantizer</code>s <ul> <li>It is an <code>nn.ModuleList</code> of <code>3</code> <code>VectorQuantizer</code>s</li> </ul> </li> <li><code>forward</code>: the main magic of this class <ul> <li>First we encode input item embedding <code>z = self.encode(x)</code></li> <li>Also init <code>residual = z</code></li> <li>Init <code>quantized_out = torch.zeros_like(z)</code> <ul> <li>The quantization output will be the sum of</li> </ul> </li> <li>Now we run a for loop through the vector quantizer layers: <ul> <li>First we compute the quantization output for this level (which contains the mapped ID for this level etc.) <ul> <li><code>vq_output: QuantizationOutput = vq_layer(residual)</code></li> </ul> </li> <li>Then we update the residual by subtracting the nearest codebook vector <ul> <li><code>residual -= vq_output.quantized.detach()</code></li> </ul> </li> <li>We accumulate the codebook vectors (with gradients) into <code>quantized_out</code> <ul> <li><code>quantized_out += vq_output.quantized_st</code></li> <li>Recall that the final representation passed to the decoder is z^=∑l=1Lel</li> </ul> </li> <li>We also accumulate the loss for each layer <ul> <li><code>vq_loss += vq_output.loss</code></li> <li>This is the codebook + commitment loss, reconstruction loss comes later</li> </ul> </li> </ul> </li> <li>Finally we get the total loss <ul> <li>Compute the reconstruction loss <ul> <li><code>x_recon = self.decode(quantized_out)</code></li> <li><code>recon_loss = F.mse_loss(x_recon, x)</code></li> <li><code>loss = recon_loss + vq_loss</code></li> </ul> </li> </ul> </li> </ul> </li> <li><code>encode_to_semantic_ids</code>: encodes an item embedding <code>x</code> to an integer tensor representing its semantic ID</li> <li><code>decode_from_semantic_ids</code>: decodes an integer tensor <code>semantic_ids</code> by looking up the codebook, summing up the levels and passing back into the <code>decoder</code></li> <li><code>kmeans_init</code> <ul> <li>Runs kmeans on one batch of embeddings to initialize the codebook vectors</li> <li>Runs kmeans to get <code>256</code> centroid vectors</li> <li>Copies these vectors into the codebook directly</li> <li>Process layer by layer</li> </ul> </li> </ul> <h3 id="finetune_qwen3_8b_vocabpy"><a class="header" href="#finetune_qwen3_8b_vocabpy">finetune_qwen3_8b_vocab.py</a></h3> This script performs Stage 1 of the qwen fine-tuning. It focuses on extending the vocabulary to include new semantic ID tokens and trains embeddings for these new tokens. <h4 id="finetuneconfig"><a class="header" href="#finetuneconfig">FineTuneConfig</a></h4> Dataclass containing config for the training <ul> <li><code>model_name</code>: <code>unsloth/Qwen3-8B</code> <ul> <li>Qn: Not instruction fine tuned?</li> </ul> </li> <li><code>load_in_4bit</code>: Set to <code>False</code> for embedding training</li> <li><code>load_in_8bit</code>: Set to <code>False</code></li> <li><code>num_proc</code>: <code>32</code></li> <li><code>enable_thinking</code>: <code>False</code> we don't need thinking mode</li> <li><code>extend_vocabulary</code>: <code>True</code></li> <li><code>codebook_levels</code>: <code>4</code></li> <li><code>codebook_size</code>: <code>256</code></li> <li><code>num_semantic_tokens</code>: <code>1024</code></li> <li><code>system_prompt</code>: see below</li> <li><code>max_training_samples</code>: <code>32000</code> limit for training embedding</li> <li><code>learning_rate</code>: <code>1e-3</code></li> <li><code>batch_size</code>: <code>32</code></li> <li><code>max_steps</code>: <code>1000</code></li> </ul> The system prompt is as follows: <blockquote> "You are a helpful AI assistant that understands and works with semantic IDs for product recommendations. Semantic IDs are hierarchical identifiers in the format <code><|sid_start|><|sid_105|><|sid_307|><|sid_705|><|sid_769|><|sid_end|></code> that encode product relationships and categories. /no_think" </blockquote> <h4 id="extend_tokenizer"><a class="header" href="#extend_tokenizer">extend_tokenizer</a></h4> <code>extend_tokenizer(model, tokenizer, config: FineTuneConfig)</code> adds semantic ID tokens to the tokenizer using Unsloth's <code>add_new_tokens</code>. <ul> <li>Note that the vocab size affects two places: <ul> <li><code>model.get_input_embeddings().weight</code>: the input embeddings</li> <li><code>model.get_output_embeddings().weight</code>: the language model head which predicts the next token</li> </ul> </li> <li>First, we make sure that the vocab size of the tokenizer matches the vocab size of both the input and output embeddings <ul> <li>We need to call <code>model.resize_token_embeddings</code> to get the model embedding sizes to match the <code>tokenizer</code></li> <li>This is because the model embeddings are padded to be multiples of <code>128</code> for CUDA optimization reasons</li> </ul> </li> <li>Next, we add new tokens using <code>unsloth.add_new_tokens</code>: <ul> <li>Special tokens of <code><|rec|></code>, <code><|sid_start|></code>, <code><|sid_end|></code></li> <li>Semantic IDs of <code><|sid_0|></code> to <code><|sid_1023|></code></li> </ul> </li> </ul> <h4 id="prepare_model"><a class="header" href="#prepare_model">prepare_model</a></h4> Prepares the model for training with some additional checks: <ul> <li>Freezes gradients for all parameters</li> <li>Unfreezes only the weights for the <code>model.get_input_embeddings()</code> and <code>model.get_output_embeddings()</code></li> <li>Checks the trainable parameter %</li> </ul> <h4 id="load_sid_dataset"><a class="header" href="#load_sid_dataset">load_sid_dataset</a></h4> Loads the semantic IDs training dataset: <ul> <li>Checks if there are texts like <code><|sid_start|></code> to make sure processing is correct</li> <li>Applies chat template to the rows (but keeps as text)</li> </ul> There are 5 distinct categories of training data: <ul> <li>SemanticID -> text: <ul> <li>Input: "Product <code><|sid_start|>...<|sid_end|></code> has title:"</li> <li>Output: "Super Mario Bros"</li> <li>Variations: ID to title, description, category, features or full context</li> </ul> </li> <li>Text -> SemanticID: <ul> <li>Input: "The product Super Mario Bros has SemanticID:"</li> <li>Output: "<code><|sid_start|>...<|sid_end|></code>"</li> <li>Variations: Similar variations to above</li> </ul> </li> <li>Sequential Recommendation: <ul> <li>Input: "Based on recent purchases etc., next item:"</li> <li>Output: "<code><|sid_start|>...<|sid_end|></code>"</li> <li>Variations: Various sequence lengths of 2, 3, or 5 items.</li> </ul> </li> <li>Semantic Understanding: <ul> <li>Input: "Products starting with <code><|sid_start|><|sid_64|> are typically:</code></li> <li>Output: "Nintendo switch games"</li> <li>Variations: Prefix to category, prefix to examples, similar items.</li> </ul> </li> <li>Multi-hop Reasoning: <ul> <li>Input: "A user who bought <code><|sid_a|></code> might also buy:"</li> <li>Output: "<code><|sid_b|></code></li> <li>Variations: Co-purchase patterns.</li> </ul> </li> </ul> <h4 id="datainspectioncallback"><a class="header" href="#datainspectioncallback">DataInspectionCallback</a></h4> Used to inspect training data and tokenization at each training step, by simply logging them to console. Patterns: <ul> <li><code>DataInspectionCallback</code> subclasses <code>transformers.TrainerCallback</code></li> <li><code>on_train_begin(self, args, state, control, **kwargs)</code>: <ul> <li>Checks the first batch of <code>train_dataloader</code></li> <li>Checks batch keys</li> <li>Check shape of <code>batch['input_ids']</code></li> <li>Check shape of <code>batch['attention_mask']</code></li> <li>Check tokens and decoded of first row etc.</li> </ul> </li> <li><code>on_log(self, args, state, control, logs=None, **kwargs)</code>: <ul> <li>Only runs if <code>state.global_step % args.logging_steps == 0</code></li> <li>Check number of SID tokens</li> <li>Decode first example and check</li> </ul> </li> </ul> <h4 id="embeddingmonitorcallback"><a class="header" href="#embeddingmonitorcallback">EmbeddingMonitorCallback</a></h4> This callback aims to check how our embeddings are shifting over time. <ul> <li>At initialization (or <code>on_train_begin</code>), we copy the state of the initial embeddings and clone detach them</li> <li>At each step, we compute the mean of the absolute difference between the current embeddings and the initial or state of embeddings from the previous step</li> <li>We also compute the per level codebook vector means etc.</li> <li>These are all logged to <code>wandb</code></li> </ul> <h4 id="semanticidgenerationcallback"><a class="header" href="#semanticidgenerationcallback">SemanticIDGenerationCallback</a></h4> This is a qualitative check to answer the question "If I ask the model for a recommendation right now, does it use the semantic ID tokens or does it just output plain text?" <ul> <li>A fixed set of test cases are used</li> <li>The test cases are <code>apply_chat_template</code>, then passed into the <code>tokenizer</code>, then <code>model.generate</code> and <code>model.decode</code></li> <li>The messages are checked whether successful (SemanticIDs generated) and success rate is tracked</li> <li>Actual completion examples are logged into wandb as well</li> </ul> <h4 id="train_embeddings"><a class="header" href="#train_embeddings">train_embeddings</a></h4> The main method. Essentially we are just using unsloth <code>SFTTrainer</code> to do the training. First, we set up <code>trl.SFTConfig</code> with a lot of the configuration we previously defined. <ul> <li>Note that <code>dataset_text_field="text"</code></li> <li><code>report_to="wandb"</code></li> </ul> The trainer <code>trl.SFTTrainer</code> is initialized with the model, tokenizer, datasets, config and callbacks. Then, <code>trainer.train()</code> is called. Note that the model and tokenizer are initialized using <code>unsloth.FastLanguageModel</code> to use unsloth's optimized triton kernels. <h3 id="finetune_qwen3_8b_fullpy"><a class="header" href="#finetune_qwen3_8b_fullpy">finetune_qwen3_8b_full.py</a></h3> The code is structurally very similar to the vocab finetuning run. The difference is that we are doing full training, so we unfreeze all parameters. Consequently, the learning rate needs to be much lower at <code>2e-5</code>. <ul> <li>Load the model from stage 1, namely <code>models/qwen3_8b_vocab_extended/final</code></li> <li>A lot of the script focuses on the callbacks to evaluate recommendation quality</li> </ul> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="cherny-2025---vibe-coding"><a class="header" href="#cherny-2025---vibe-coding">Cherny 2025 - Vibe Coding</a></h1> Boris Cherny (creator of Claude code) has a <a href="https://x.com/bcherny/status/2007179833990885678">series of tweets</a> on how he uses Claude. This should be a good point of reference. <h1 id="1-run-claudes-in-5-terminals"><a class="header" href="#1-run-claudes-in-5-terminals">1. Run Claudes in 5 Terminals</a></h1> Tabs are numbered 1-5. Configure system notifications to know when a claude needs input. The <a href="https://code.claude.com/docs/en/terminal-config#iterm-2-system-notifications">link</a> provided is for iterm2 on Mac but probably can do something similar for linux or WSL. <h1 id="2-run-claudes-online"><a class="header" href="#2-run-claudes-online">2. Run Claudes online</a></h1> Concurrently to the local terminals, he has 5-10 claudes running online at <a href="https://claude.ai/code">https://claude.ai/code</a>. These can be triggered either by handing off terminal sessions to the web (using <code>&</code>) or manually kicking off sessions in the browser. <ul> <li>Question: Does claude run in its own instance when running online? Doesn't that mean the context it has is different from the local context?</li> </ul> <h1 id="3-use-opus-45-w-thinking-all-the-time"><a class="header" href="#3-use-opus-45-w-thinking-all-the-time">3. Use Opus 4.5 w/ Thinking all the time</a></h1> Basically, use the biggest and slowest model all the time. Although it runs slower, it will complete things more reliably and with things running async, the reliability wins out in the end. <h1 id="4-use-a-claudemd"><a class="header" href="#4-use-a-claudemd">4. Use a CLAUDE.md</a></h1> The claude code repo has a single <code>CLAUDE.md</code> that is checked into git and everyone contributes to. Everytime claude fails on a task, or does something in a slightly undesirable way (think of all the <code>nit</code> comments), we can add a line to <code>CLAUDE.md</code> so that it does not do that again. <h1 id="5-use-claude-to-add-stuff-to-claudemd"><a class="header" href="#5-use-claude-to-add-stuff-to-claudemd">5. Use <code>@.claude</code> to add stuff to CLAUDE.md</a></h1> Using the claude code github action <code>/install-github-action</code>, tagging @.claude in PR comments we can instruct claude to add stuff to CLAUDE.md. e.g. <code>@.claude add to CLAUDE.md to always use string literals instead of enums</code>. <h1 id="6-start-sessions-in-plan-mode"><a class="header" href="#6-start-sessions-in-plan-mode">6. Start sessions in plan mode</a></h1> Start in plan mode <code>shift+tab twice</code>, and go back and forth. Once the plan is good, switch to auto accept edits and Claude usually 1-shots the task. A good plan is super important. <h1 id="7-use-slash-commands-for-frequent-workflows"><a class="header" href="#7-use-slash-commands-for-frequent-workflows">7. Use slash commands for frequent workflows</a></h1> Commands are commonly executed workflows that live inside <code>.claude/commands/</code>. These can be invoked in claude chat e.g. <code>/commit-push-pr</code>. This command computes git status, and some other info before getting claude to summarize and create a PR. More details at <a href="https://code.claude.com/docs/en/slash-commands#bash-command-execution">https://code.claude.com/docs/en/slash-commands#bash-command-execution</a>. The <code>/commit-push-pr</code> command may be viewed <a href="https://github.com/anthropics/claude-code/blob/main/.claude/commands/commit-push-pr.md">here</a> <h1 id="8-use-sub-agents-for-more-involved-workflows"><a class="header" href="#8-use-sub-agents-for-more-involved-workflows">8. Use sub-agents for more involved workflows</a></h1> Slash commands are for short, simple workflows. Sub agents (living in <code>.claude/agents</code>) are also frequent workflows but more involved. For e.g., code simplification to simplify code after claude is done working. The code simplifier agent may be viewed <a href="https://github.com/anthropics/claude-plugins-official/blob/main/plugins/code-simplifier/agents/code-simplifier.md">here</a> <h1 id="9-use-a-posttooluse-hook"><a class="header" href="#9-use-a-posttooluse-hook">9. Use a PostToolUse hook</a></h1> Use a <code>PostToolUse</code> hook to format code after claude is done. Basically just <code>bun format</code>. We can also do <code>ruff format</code> etc. depending on the language. <h1 id="10-configure-permissions"><a class="header" href="#10-configure-permissions">10. Configure permissions</a></h1> Do not use <code>--dangerously-skip-permissions</code>. Instead use <code>/permissions</code> and add commands that we know are safe in our environment. <h1 id="11-tool-use"><a class="header" href="#11-tool-use">11. Tool Use</a></h1> Use <code>.mcp.json</code> to specify the MCP servers we allow claude to connect to and use tools. <h1 id="12-long-running-tasks"><a class="header" href="#12-long-running-tasks">12. Long Running Tasks</a></h1> Use <code>--permission-mode=dontAsk</code> or <code>--dangerously-skip-permissions</code> so that claude will not get stuck. Also use hooks to ask the agent to periodically verify what is going on for monitoring. <h1 id="13-invest-in-the-verification-loop"><a class="header" href="#13-invest-in-the-verification-loop">13. Invest in the Verification Loop</a></h1> This is the most important step. Claude needs a way to verify its work, whether it is through some unit tests, integration tests or browser usage etc. Having a verification feedback loop will 2x or 3x the quality. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="experiments-6"><a class="header" href="#experiments-6">Experiments</a></h1> Collection of my personal experiments to validate some hypotheses. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="rq-vae-fast-decoding"><a class="header" href="#rq-vae-fast-decoding">RQ-VAE Fast Decoding</a></h1> <a href="https://github.com/charleslow/rq_vae">Github Repo: rq_vae</a> Experiments on applying RQ-VAE to fast text decoding in latent space. <h2 id="main-idea-2"><a class="header" href="#main-idea-2">Main Idea</a></h2> We use ideas from two papers: <ul> <li><a href="https://charleslow.github.io/notebook/book/papers/kaiser_2018.html">Kaiser 2018: Fast Decoding in Sequence Models using Discrete Latent Variables</a> <ul> <li>Applies the idea of representing tokens in a shorter latent space, and then doing autoregressive text translation in the latent space, then upsample back to token space</li> <li>Still uses old VQ-VAE discretization which has issues</li> </ul> </li> <li><a href="https://charleslow.github.io/notebook/book/papers/lee_2022.html">Lee 2022: Autoregressive Image Generation using Residual Quantization</a> <ul> <li>Better way of doing discretization, using a codebook with multiple levels instead of a flat codebook</li> <li>Some tricks of using a specialized transformer for decoding that is faster</li> </ul> </li> </ul> <h2 id="experiment-goal"><a class="header" href="#experiment-goal">Experiment Goal</a></h2> Verify if we can achieve significant speedup by decoding in latent space without losing accuracy <ul> <li>Check perplexity on held out data</li> <li>Check codebook usage</li> <li>Compare against various base models (including qwen0.6B) on standard LLM tasks</li> </ul> Other goals: <ul> <li>Test if RQ-transformer matches standard transformer decoding for the same compute</li> <li>Test if using a pretrained decoder backbone is necessary</li> <li>Test performance - inference speed trade-off as we increase compression factor</li> <li>Test codebook vocabulary size to codebook depth tradeoff</li> <li>Test if unsloth speeds up inference significantly</li> <li>Test if we can scale to 8B models</li> </ul> <h2 id="overall-flow"><a class="header" href="#overall-flow">Overall Flow</a></h2> <ul> <li>Start with text input (e.g. <code>fineweb</code>)</li> <li>Encode into a latent representation that is a shorter sequence <ul> <li>Run through pretrained backbone (e.g. <code>qwen0.6B</code>)</li> <li>Downsample using convolutions</li> </ul> </li> <li>Run RQ-VAE on the latent representation to discretize into latent tokens <ul> <li>This is the fast decoding part</li> </ul> </li> <li>Decode latent tokens back into original textual space <ul> <li>Upsample using de-convolutions</li> <li>Run through pretrained backbone (e.g. <code>qwen0.6B</code>)</li> </ul> </li> </ul> <h2 id="training-sequence"><a class="header" href="#training-sequence">Training Sequence</a></h2> The trainable components are encoder, RQ-codebook, decoder for the RQ-VAE portion, and also the RQ-transformer for the latent decoding. <ul> <li>Train the RQ-VAE jointly <ul> <li>Freeze pretrained qwen backbone for warmup</li> <li>Then unfreeze for full training</li> </ul> </li> <li>Train the RQ-transformer on latent sequences</li> </ul> <h2 id="inference-sequence"><a class="header" href="#inference-sequence">Inference Sequence</a></h2> At inference time, the goal is to do autoregressive text completion (like normal LLMs): <ul> <li>Given input text</li> <li>Use encoder to encode into latent space</li> <li>Use RQ-transformer to auto-regressively predict full latent sequence</li> <li>Use decoder to decode back into textual space</li> </ul> <h2 id="code-deep-dive"><a class="header" href="#code-deep-dive">Code Deep Dive</a></h2> Here we dive into dissecting the code in detail. The code structure is: <pre><code>│ model/ │ ├── rq_vae.py │ │ └── RQVAE ─────────────┬─► TextEncoder │ │ ├─► ResidualQuantizer │ │ └─► TextDecoder │ ├── encoder.py │ ├── decoder.py │ ├── quantizer.py │ ├── rq_transformer.py │ │ ├── SpatialTransformer │ │ ├── DepthTransformer │ │ ├── RoPEAttention │ │ └── TransformerBlock │ └── layers.py │ └── SwiGLUTransformerLayer │ train_vae.py │ │ ├── RQVAELightningModule │ │ └── RQVAEDataModule │ train_transformer.py │ ├── RQTransformerLightningModule │ └── RQTransformerDataModule </code></pre> <h3 id="encoderpy"><a class="header" href="#encoderpy">encoder.py</a></h3> Implements <code>TextEncoder</code> which encodes an input text into latent embeddings representation (shorter sequence). We focus on the forward pass: <ul> <li>Input text sequence: <code>batch_size, seq_len</code></li> <li>Encode using pretrained <code>qwen</code> backbone: <code>batch_size, seq_len, hidden_size</code></li> <li>Downsample into shorter latent sequence: <code>batch_size, compressed_len, hidden_size</code> <ul> <li>Use strided convolutions to downsample the sequence length (halves length each step)</li> <li>More details below</li> </ul> </li> <li>Linear projection into <code>batch_size, seq_len, latent_dim</code> <ul> <li>The <code>latent_dim</code> is the dimension of our RQ-VAE codebook</li> </ul> </li> <li>Refine latents using self-attention: <code>batch_size, seq_len, latent_dim</code> <ul> <li>Do it with <code>SwiGLUTransformerLayer</code> for <code>num_latent_layers</code> times</li> </ul> </li> <li>Output the latent representation</li> </ul> <blockquote> Question: How much information do we lose in the convolutional downsampling? How can we encourage the latent space to store more information? </blockquote> Diving a bit more into the convolutional downsampling process: <ul> <li><code>nn.Conv1d</code> applies a convolutional filter which takes a weighted sum of the input values within the window <ul> <li><code>kernel_size</code> is the size of the convolution window / filter. Larger value means we average more values together.</li> <li><code>stride</code> is how many steps we move when sliding the window. <code>stride=2</code> approximately halves the sequence length.</li> <li><code>padding</code> is the number of zero-padded cells we add to each side of the input</li> </ul> </li> <li>For an input sequence of Lin, padding p, kernel_size k, and stride s, the output length is: ⌊Lout=sLin+2p−k+1⌋</li> <li>The way to think about the formula is that we start with the first window (hence +1 at the end) and then count the number of strides to take <ul> <li>First we add padding 2p to pad the input size</li> <li>Subtract k because we already "occupy" the first kernel width</li> <li>Add a <code>floor</code> operator to account for partial strides which we discard</li> </ul> </li> <li>Can use the <a href="experiments/./conv_viz.html">interactive convolution explorer</a> to visualize the convolution window</li> </ul> <h3 id="decoderpy"><a class="header" href="#decoderpy">decoder.py</a></h3> The decoder <code>TextDecoder</code> is very similar to the encoder, except we go in the other direction. From a latent sequence, we upsample using de-convolutions to return to the textual space. Again we focus on the forward pass: <ul> <li>Input is a sequence of compressed latents: <code>batch_size, compressed_len, latent_dim</code></li> <li>Refine latents with self-attention: <code>batch_size, compressed_len, latent_dim</code></li> <li>Linear projection to hidden size: <code>batch_size, compressed_len, hidden_size</code></li> <li>Upsample via transposed convolutions: <code>batch_size, seq_len, hidden_size</code></li> <li>Each layer 2x expands sequence length</li> <li>Process through Qwen3 backbone (one-shot, not autoregressive): <code>batch_size, seq_len, hidden_size</code></li> <li>Linear projection to vocabulary logits: <code>batch_size, seq_len, vocab_size</code></li> </ul> <h3 id="quantizerpy"><a class="header" href="#quantizerpy">quantizer.py</a></h3> Now we dive into the residual quantization portion. The Quantizer class represents one level of the codebook. It is initialized with parameters: <ul> <li><code>dim</code>: latent dimension of codebook vectors</li> <li><code>codebook_size</code>: number of vectors for this level of the codebook</li> <li><code>ema_decay</code>: the rate of updating codebook vectors (analogous to <code>1 - learning_rate</code>)</li> <li><code>threshold_ema_dead_code</code>: number of dead codes we allow before we re-assign these codebook vectors</li> </ul> We register these buffers at init: <ul> <li><code>ema_cluster_size</code>: <code>torch.zeros(codebook_size)</code> <ul> <li>Tracks the number of times each codebook vector is used in an exponential moving average</li> </ul> </li> <li><code>ema_embed_sum</code>: <code>torch.zeros(codebook_size, dim)</code> <ul> <li>Tracks the sum of encoder embeddings assigned in an exponential moving average</li> </ul> </li> </ul> The <code>quantize</code> method assigns an input encoder embedding <code>x</code> to the nearest codebook vector. Note that <code>x</code> would represent a <code>residual</code> vector (depending on the level we are at in the codebook) in the RQ-VAE architecture. <ul> <li>First we use <code>torch.cdist</code> to find distances to all the codebook vectors</li> <li>Then we use <code>torch.argmin</code> to get the index of the nearest codebook vector</li> <li>Then we retrieve the codebook embedding corresponding to this index</li> <li>We return a tuple <code>(quantized, indices)</code></li> </ul> The <code>update_codebook_ema</code> acts like a gradient update to the codebook vectors (following <a href="experiments/../papers/van_den_oord_2017.html">Oord 2017</a>). <ul> <li><code>one_hot: batch_size, codebook_size = F.one_hot(indices, codebook_size)</code> <ul> <li>After we run <code>quantize</code> we convert the nearest neighbour <code>indices</code> into a one-hot matrix</li> </ul> </li> <li><code>cluster_size = one_hot.sum(dim=0)</code> <ul> <li>The one-hot matrix is summed on dimension <code>0</code> to give the number of counts per cluster</li> <li>Shape is <code>codebook_size</code></li> </ul> </li> <li><code>embed_sum = one_hot.t() @ x</code> <ul> <li>This is the sum of encoder embeddings based on assignment to clusters</li> <li>Shape is <code>codebook_size, latent_dim</code></li> </ul> </li> <li>The exponential moving average of <code>cluster_size</code> and <code>embed_sum</code> are taken and assigned to <code>self.ema_cluster_size</code> and <code>self.ema_embed_sum</code> respectively <ul> <li>The codebook vectors are updated as <code>self.ema_embed_sum / self.ema_cluster_size</code></li> </ul> </li> <li>Dead code checking to avoid codebook collapse <ul> <li>Codebook collapse is the case where some codebook vectors are so far away from the encoder distribution that they never get used</li> <li>We detect these codes by tracking their exponential moving average of counts</li> <li>If the average assignment falls below a certain threshold, we delete these codebook vectors and re-initialize them to an average of a few random encoder embeddings in the batch</li> </ul> </li> </ul> There is also <code>compute_perplexity</code>, which is used to measure the distribution of codebook utilisation using perplexity <ul> <li>Perplexity is defined as exp(entropy)=exp(∑i=1K−pilogpi) of a categorical distribution with K categories</li> <li>We know that entropy ranges from 0 to logK</li> <li>So perplexity ranges from 1 to K</li> <li>We want the perplexity to be close to K, showing that codebook utilisation is close to uniform</li> </ul> <blockquote> Question: <ul> <li>Is uniform codebook utilisation optimal, or is there some way to reason about the ideal distribution of codebook utilisation?</li> <li>There seems to be some reason to reduce the codebook size as we go into deeper levels, since the variance decreases and there is higher risk of modelling noise. Is there some information theoretic way to dynamically adjust the codebook size based on the amount of information gain from that codebook level?</li> </ul> </blockquote> The ResidualQuantizer class is a stack of Quantizers based on the desired number of <code>codebook_levels</code>. The <code>forward</code> pass is where most of the logic resides: <ul> <li>We receive an input tensor <code>x</code> with shape <code>batch_size, seq_len, latent_dim</code></li> <li>We run a for loop through all the <code>codebook_levels</code>: <ul> <li><code>quantized, indices = quantizer.quantize(residual)</code> <ul> <li>Run the quantizer to get the assigned codebook vector and indices</li> </ul> </li> <li>Update the exponential moving average for the quantizer</li> <li><code>commitment_loss = F.mse_loss(residual, quantized.detach())</code> <ul> <li>Compute the commitment loss that pushes encoder output to be near codebook vectors</li> <li>Notice the stop gradient on the codebook vectors</li> </ul> </li> <li>Compute perplexity and keep track</li> <li>Accumulate <code>quantized</code> into <code>quantized_out</code> <ul> <li>Recall that RQ-VAE represents each token as a sum of the codebook vectors across all levels</li> <li>Hence we sum <code>quantized</code> at each level to get our final representation</li> <li>The straight-through estimator is applied to <code>quantized_out</code> to get gradients flowing back to the encoder</li> </ul> </li> </ul> </li> </ul> <h3 id="rqvaepy"><a class="header" href="#rqvaepy">rqvae.py</a></h3> This script combines the encoder, decoder and residual quantizer into one module for joint training. The main method is <code>forward</code>, which does the <code>encode -> quantize -> decode</code> pass and then computes loss to train all its components. <ul> <li> Firstly we receive <code>input_ids</code> and <code>attention_mask</code> with shape <code>batch_size, seq_len</code> from the tokenizer </li> <li> We call <code>self.encode</code> which does two things and returns a <code>dict</code>: <ul> <li><code>latent = self.encoder(input_ids, attention_mask)</code></li> <li><code>quant_out = self.quantizer(latent)</code></li> <li>Returns: <ul> <li><code>latent</code>: the encoder representation in latent space</li> <li><code>quantized</code>: the quantized representation (sum of codebook vectors for each position)</li> <li><code>indices</code>: the latent token sequence</li> <li><code>commitment_loss</code></li> <li><code>perplexities</code>: for tracking purpose</li> <li><code>dead_code_replacements</code>: for tracking purpose</li> </ul> </li> </ul> </li> <li> <code>logits = self.decoder(quantized, seq_len)</code> <ul> <li>This decodes the quantized representation back into logits for each position</li> <li>The logits represent our predictions for each token probability for each position</li> </ul> </li> <li> The logits are used to compute the <code>reconstruction_loss</code> <ul> <li>The <code>reconstruction_loss</code> is the cross entropy between predicted logits and actual labels</li> <li>The actual labels default to <code>input_ids</code> if labels are not explicitly provided</li> </ul> </li> <li> The total loss is computed as <code>reconstruction_loss + commitment_loss</code> </li> </ul> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="nlp-course"><a class="header" href="#nlp-course">NLP Course</a></h1> Based on lectures from Graham Neubig's <a href="https://www.youtube.com/watch?v=6NeTO61qc4M">Advanced NLP Class</a> in 2024. <h1 id="lecture-1-introduction"><a class="header" href="#lecture-1-introduction">Lecture 1: Introduction</a></h1> <a href="https://www.youtube.com/watch?v=6NeTO61qc4M">Lecture 1 Link</a> A general framework for NLP systems: create a function to map an input X into an output Y, where X and/or Y involved language. Tasks include: <ul> <li>Language Modelling</li> <li>Translation</li> <li>Text Classification</li> <li>Langugage Analysis</li> <li>Image Captioning</li> </ul> Challenges of NLP <ul> <li>Low frequency words</li> <li>Conjugation (e.g. adjectives modifying a word)</li> <li>Negation</li> <li>Metaphor, analogy</li> <li>Other languages</li> </ul> Topics: <ol> <li>Language Modelling Fundamentals</li> <li>Training and Inference Methods</li> <li>Experimental Design and Evaluation</li> <li>Advanced Training and Architectures</li> <li>NLP Applications</li> <li>Linguistics and Multi-Linguality</li> </ol> <h1 id="lecture-2-word-representation-and-text-classification"><a class="header" href="#lecture-2-word-representation-and-text-classification">Lecture 2: Word Representation and Text Classification</a></h1> Subword models aim to address issue of low-frequency words. How many words are there in the English language? This is a bit of a trick question because do we consider <code>company</code> or <code>companies</code> different words? By using subwords we can greatly reduce the vocabulary size to around <code>60k</code> which is used in modern tokenizers. Another way to address this is to use character-level models. The main issue of character-level models is that our sequences become very long, so for a fixed sequence length we cover a lot less text in the same batch size. Using subwords saves compute and memory. Byte Pair Encoding (2015) is a very simple method to create subwords. It simply incrementally combines the most frequent token pairs together, starting at the character level. e.g. starting with the sentence <code>newest, widest</code>, we first combine <code>es</code> into a token, then <code>est</code>, and so on. Another way to do this is to use Unigram Models (e.g. Kudo 2018). First we use a unigram LM that generates all words in the sequence independently. We then pick a vocabulary that maximizes the log likelihood of the corpus given a fixed vocabulary size. The optimization process is using the EM algorithm. <code>Sentencepiece</code> is a highly optimized library to train both these types of subword models. Subword nuances: <ul> <li>Subword models are hard to use multilingually because they will over-segment less common languages naively</li> <li>Subword segments are sometimes arbitrary (e.g. is it <code>es t</code> or <code>e st</code>)?</li> </ul> Pytorch vs Tensorflow / JAX <ul> <li>Pytorch is more widely used</li> <li>Pytorch favour dynamic execution vs compile + execute</li> </ul> Training transformers typically uses a learning rate schedule. Start low, increase in the middle and decrease toward the end. Starting with a high learning rate will lead to weird models because transformers are sensitive. <h1 id="lecture-3-language-modelling"><a class="header" href="#lecture-3-language-modelling">Lecture 3: Language Modelling</a></h1> <a href="https://www.youtube.com/watch?v=69EAJOwV3Es">Lecture 3 Link</a> Generative vs Discriminative models: <ul> <li>Discriminative model: a model that calculates the probability of a latent trait given the data, i.e. P(Y∣X), where Y is say the sentiment and X is the text</li> <li>Generative model: a model that calculates the probability of the data itself i.e. P(X)</li> </ul> Generative models for text are called language models. We can use LMs to generate sentences. We can also use them to score sentences (e.g. when a sentence is a concatenation of a question and answer pair). Can also use LMs to correct sentences. Auto-regressive language models compute P(X)=∏i=1IP(xi∣x1,...,xi−1). Question: why do we do it auto-regressively instead of computing the entire sequence x at once? The problem is computational - predicting the next token has a space of ∣V∣∼60,000, but predicting the entire sequence is in the order of ∣V∣N where N is the length of the sequence, which is currently untractable. That being said, if we can model the entire sequence, it will probably be a lot more efficient than the auto-regressive way. The simplest language model is count-based unigram model. By making an indepenence assumption P(xi∣x1,...,xi−1)∼P(xi), we just ignore all the previous words and predict the probability of a word occurring. The maximum likelihood estimation for the probability of a given token xi will simply be: PMLE(xi)=∑x<svg xmlns="http://www.w3.org/2000/svg" width='100%' height='0.26em' viewBox='0 0 600 260' preserveAspectRatio='none'><path d='M200 55.538c-77 0-168 73.953-177 73.953-3 0-7 -2.175-9-5.437L2 97c-1-2-2-4-2-6 0-4 2-7 5-9l20-12C116 12 171 0 207 0c86 0 114 68 191 68 78 0 168-68 177-68 4 0 7 2 9 5l12 19c1 2.175 2 4.35 2 6.525 0 4.35-2 7.613-5 9.788l-19 13.05c-92 63.077-116.937 75.308-183 76.128 -68.267.847-113-73.952-191-73.952z'/></svg>(x<svg xmlns="http://www.w3.org/2000/svg" width='100%' height='0.26em' viewBox='0 0 600 260' preserveAspectRatio='none'><path d='M200 55.538c-77 0-168 73.953-177 73.953-3 0-7 -2.175-9-5.437L2 97c-1-2-2-4-2-6 0-4 2-7 5-9l20-12C116 12 171 0 207 0c86 0 114 68 191 68 78 0 168-68 177-68 4 0 7 2 9 5l12 19c1 2.175 2 4.35 2 6.525 0 4.35-2 7.613-5 9.788l-19 13.05c-92 63.077-116.937 75.308-183 76.128 -68.267.847-113-73.952-191-73.952z'/></svg>)ctrain(xi) Detail: parametrizing in log space. Multiplication of probabilities are re-expressed as additions in log space, because otherwise, if we multiply 100 probabilities together for a sequence of 100 tokens, we will easily underflow the numeric precision. P(X)=x=1∏NP(xi)logP(X)=x=1∑NlogP(xi) Correspondingly, we can define the parameters θxi=logP(xi). Moving on to higher order n-gram models, the idea is to limit the context length to one-word before the token we are predicting, and then count: PML(xi∣Xi−n+1,...,xi−1):=c(xi−n+1,...,xi−1)c(xi−n+1,...,xi) e.g. P(example | this is an) = c(this is an example) / c(this is an). Due to sparsity of data, we need to add smoothing to deal with zero counts, i.e. instead of just using tri-gram, we smooth tri-gram and bi-gram probabilities together: P(xi∣xi−n+1,...,xi−1)=λPML(xi∣xi−n+1,...,xi−1)+(1−λ)P(xi∣xi−n+2,...,xi−1) e.g. P(example | this is an)=λPML(example | this is an)+(1−λ)PML(example | is an). More sophisticated smoothing techniques are studied in <a href="https://dash.harvard.edu/bitstream/handle/1/25104739/tr-10-98.pdf">Goodman 1998: An Empirical Study of Smoothing Techniques for Language Modelling</a>. Problems: <ol> <li>Cannot share strength amongst similar words, e.g. <code>car</code> and <code>bicycle</code></li> <li>Cannot condition on context with intervening words, e.g. <code>Dr Jane Smith</code> vs <code>Dr Gertrude Smith</code></li> <li>Cannot handle long-distance dependencies, e.g.<code>tennis</code> and <code>racquet</code> in <code>for tennis class he wanted to buy his own racquet</code></li> </ol> The standard toolkit for n-gram models is <a href="https://github.com/kpu/kenlm">kenlm</a> which is extremely fast and scalable, written in c++. <h2 id="evaluating-language-models"><a class="header" href="#evaluating-language-models">Evaluating Language Models</a></h2> <ol> <li>Log Likelihood:</li> </ol> LL(Xtest)=X∈Xtest∑logP(X) <ol start="2"> <li>Per-word Log Likelihood:</li> </ol> WLL(Xtest)=∑X∈Xtest∣X∣1X∈Xtest∑logP(X) <ol start="3"> <li>Per-word Cross Entropy:</li> </ol> H(Xtest)=−∑X∈Xtest∣X∣1X∈Xtest∑log2P(X) Aside: Any probabilistic distribution can also be used to compress data. The entropy measure is closely related to the number of bits needed to store the data based on our language model. <ol start="4"> <li>Perplexity. Lower is better. Perplexity is the number of times we need to sample from the probability distribution until we get the answer right.</li> </ol> PPL(Xtest)=2H(Xtest)=e−WLL(Xtest) <h2 id="other-desiderata-of-lms"><a class="header" href="#other-desiderata-of-lms">Other Desiderata of LMs</a></h2> Calibration Guo 2017. Formally, we want the model probability of the answer matching the actual probability of getting it right. Typically we measure calibration by bucketing the model output probabilities and calculating expected calibration error: ECE=m=1∑Mn∣Bm∣∣acc(Bm)−confidence(Bm)∣ where m represents a sub-segment of the data which corresponds to a confidence interval from the model. How do we calculate answer probabilities? e.g. <code>the university is CMU</code> and <code>the university is Carnegie Mellon University</code> should both be acceptable. <ul> <li>One way is to use paraphrases to substitute phrases <code>Jiang 2021</code>.</li> <li>One way is to just ask the model to generate a confidence score, see <a href="https://aclanthology.org/2023.emnlp-main.330.pdf">Tian 2023 - Just ask for calibration</a></li> </ul> Another desirable characteristic is efficiency. Some metrics are: <ul> <li>Memory usage (load model only, peak memory usage)</li> <li>Latency (to first token, to last token)</li> <li>Throughput</li> </ul> Some efficiency tips: <ul> <li>On modern hardware doing 10 operations of size 1 is much slower than doing 1 operation of size 10</li> <li>CPUs are like motorcycles and GPUs are like airplanes</li> <li>Try to avoid memory moves between CPU and GPU, and if we need to move memory, do so as early as possible (as GPU operations are asynchronous).</li> </ul> <h1 id="lecture-4-sequence-modelling"><a class="header" href="#lecture-4-sequence-modelling">Lecture 4: Sequence Modelling</a></h1> <a href="https://www.youtube.com/watch?v=x3U2zVhrgJ8">Lecture 4 Link</a> NLP is full of sequential data, especially those containing long range dependencies. References can also be complicated, e.g. <code>the trophy would not fit in the suitcase because it was too big</code>. What does <code>it</code> refer to? These are called winograd schemas. Types of tasks: <ul> <li>Binary classification</li> <li>Multi-class classification</li> <li>Structured Prediction. e.g. predicting the parts-of-speech tag for each word in the sentence</li> </ul> Sequence labelling is an important category of tasks in NLP: <ul> <li>e.g. Parts of speech tagging</li> <li>Lemmatization</li> <li>Morphological tagging, e.g. <code>PronType=prs</code></li> </ul> Span labelling: <ul> <li>Named entity resolution</li> <li>Syntactic Chunking</li> <li>Entity Linking</li> <li>Semantic role labelling</li> </ul> We can treat span labelling as a sequence labelling task by using <code>beginning</code>, <code>in</code> and <code>out</code> tags. Three major types of sequence modelling: <ol> <li>Recurrence</li> <li>Convolutional</li> <li>Attentional</li> </ol> Recurrent neural networks essentially unroll a computational graph through "time" (where time is the position in the sequence). This results in a vanishing gradient problem, where the gradient on the last token is largest, and the gradient becomes smaller as we move backwards towards the first token. This problem is not just present in RNNs, but in general for computation graphs, if there is important information, adding direct connections from the important nodes to the loss is one way to improve performance. This is the motivation for residual or skip-connections. LSTM is one way of solving this problem - the basic idea is to make additive connections between time steps, which does not result in vanishing. The idea is to have different "gates" that control the information flow, and use additive connections. Additive connections solves the vanishing gradient problem because it does not modify the gradient (?). Attention score functions: <ul> <li> Original paper in Bahdanau 2015 used a multi-layer perceptron for the attention function, which is very expressive but has extra parameters a(q,k)=w2Ttanh(W1[q;k]) </li> <li> Bilinear in Luong 2015: a(q,k)=qTWk </li> <li> Dot product: a(q,k)=qTk </li> <li> Scaled Dot product (Vaswani 2017): a(q,k)=∣k∣<svg xmlns="http://www.w3.org/2000/svg" width='400em' height='1.28em' viewBox='0 0 400000 1296' preserveAspectRatio='xMinYMin slice'><path d='M263,681c0.7,0,18,39.7,52,119 c34,79.3,68.167,158.7,102.5,238c34.3,79.3,51.8,119.3,52.5,120 c340,-704.7,510.7,-1060.3,512,-1067 l0 -0 c4.7,-7.3,11,-11,19,-11 H40000v40H1012.3 s-271.3,567,-271.3,567c-38.7,80.7,-84,175,-136,283c-52,108,-89.167,185.3,-111.5,232 c-22.3,46.7,-33.8,70.3,-34.5,71c-4.7,4.7,-12.3,7,-23,7s-12,-1,-12,-1 s-109,-253,-109,-253c-72.7,-168,-109.3,-252,-110,-252c-10.7,8,-22,16.7,-34,26 c-22,17.3,-33.3,26,-34,26s-26,-26,-26,-26s76,-59,76,-59s76,-60,76,-60z M1001 80h400000v40h-400000z'/></svg>qTk </li> </ul> Note that the attention mechanism in Vaswani 2017 is essentially a bilinear function + scaling, because the weight matrices are involved in obtaining the query and key vectors. Comparing RNN, convolution and attention, for a token at position <code>200</code> to attend to token at position <code>1</code>: <ul> <li>RNN will take 199 steps through the computation graph</li> <li>Convolution will take 20 steps if the convolution window is 10</li> <li>Attention will take 1 step</li> </ul> This is an advantage of the attention representation of text. Another advantage of attention is the computation speed. For a given text of <code>N</code> tokens, in order to produce <code>N-1</code> token predictions for the next word at each step: <ul> <li>RNN will have to sequentially run <code>N-1</code> steps</li> <li>Attention will do it in 1 step by masking the attention matrix suitably. This makes it much faster to train attention models.</li> </ul> Truncated Backpropagation is one technique to reduce computation. e.g. in the context of RNN, we might forward propagate through positions <code>1-10</code> and compute the gradients. For the next positions <code>11-20</code>, the hidden state from position <code>10</code> is used for forward propagation, but we do not backpropagate the gradients back to positions <code>1-10</code>. <h1 id="lecture-5-transformers"><a class="header" href="#lecture-5-transformers">Lecture 5: Transformers</a></h1> <a href="https://www.youtube.com/watch?v=QkGwxtALTLU">Lecture 5 Link</a> Transformers started as cross attention (Bahdanau 2014) and evolved to self-attention (Vaswani 2017). There are two main types of transformers: <ol> <li> Encoder-Decoder model, e.g. T5 or BART. Encoder-decoder models have a separate encoder and decoder module. The encoder takes the input sequence and passes it into a transformer block to return a "context" embedding. The decoder then takes in both the "context" embedding and the output sequence (e.g. the same sentence in French for translation, or the same sentence itself shifted right for self-attention) with appropriate masking and generates a probability for each position in the output sequence. Note that the encoder module has un-masked full attention over the entire input sequence, but the decoder module is masked. The benefit of the encoder-decoder model is that we get an embedding representation of a given input sequence which can be used for downstream tasks. But it is also less flexible because we need to have a concept of the "input sequence" vs the "output sequence", which is suitable for tasks like translation but not in text generation tasks there is no clear input/output. </li> <li> Decoder only model, e.g. GPT or LLaMa. The decoder model simply takes in the input sequence and passes it through a transformer module, resulting in a probability for each position in the input sequence. Using appropriate causal masking, we can train the network to predict the next token at each position given only information about tokens at previous positions. The benefit of the decoder model is fewer parameters than the encoder-decoder model. It's also more suitable for text generation tasks. In the encoder decoder framework, at each time step we need to recompute the encoder representation, because the representation at earlier positions can change with the addition of a new token. In the decoder only framework, the previous cached Q, K, V values do not change due to the causal masking, so we can re-use those in decoding the next time step. We should thus expect decoder only models to be faster at decoding. </li> </ol> <h2 id="core-concepts"><a class="header" href="#core-concepts">Core concepts</a></h2> Multi-head attention. The intuition for having multiple attention heads is that information from different parts of the sentence can be useful to disambiguate in different ways. Typically for a given word to attend to nearby words is useful for learning syntax, but attending to further words is useful for learning semantics. Multi-head attention basically comprises of multiple attention modules, each with their own weights WQi,WKi,WVi. The attention outputs are computed for each head, and then concatenated together. Since the resulting matrix will have nhead times the size on one dimension, we pass it through a final linear transformation to get the desired dimension size (as though we only used one attention head). In practice the attention weights across all the heads are concatenated together first before the matrix multiplication to vectorize the computation. The resulting matrix is then sliced and attention computed on each head, before concatenating together again for the final matrix multiplication. Positional Encoding. There is no notion of position in the transformer. Positional encoding simply adds an embedding at each position to the word embedding to encode this information. Note that it is added right at the beginning to the raw token embeddings. <ol> <li> Sinusoidal Encoding was the original proposal in the Vaswani 2017 paper. The position embedding is a fixed vector at each position t, where the ith element is sin(ωk⋅t) for even i=2k and cos(ωk⋅t) for odd i=2k+1, and ωk:=100002k/d1. Here k is an index on the dimension going from 1,...2d and d is the dimension of the embedding. The method is rather counter-intuitive but the basic idea is that we want the dot product between two embeddings to be high when the relative position is near and decay as we move away. The intuition is covered nicely in <a href="https://kazemnejad.com/blog/transformer_architecture_positional_encoding/">Kazemnejad 2019</a> and basically we can think of each positional embedding as analogous to a binary encoding represented by 0,1 bits at each dimension. Instead of a hard 0 or 1, the sine and cosine functions provide a smoothed version so that we get a nice relative decay. Note that this method does allow us to extrapolate the position embedding to longer sequences that we have not seen before in training. </li> <li> Learnable Embeddings. This was proposed in Shaw 2018, which basically just allows the embedding at each position to be a learnable vector. The problem of this approach is that it becomes impossible to extrapolate to longer sequences in inference time. </li> <li> Rotary Positional Encodings or RoPE. This was proposed in Su 2021. The fundamental idea is that we want the dot product of embeddings to result in a function of relative position. Specifically, we desire that for given positions m,n, and the respective word embeddings xm,xn, the dot product of the resulting embeddings may be expressed purely as a function of their relative distance m−n, and in so doing we lose notion of the absolute position entirely. fq(xm,m)⋅fk(xn,n)=g(xm,xn,m−n) The paper uses trigonometry and imaginary numbers to come up with a function that satisfies this property. The benefit of losing notion of absolute position entirely means that we can extrapolate to longer sequences that we have not seen before, and RoPE extrapolates better than sinusoidal embeddings. LLaMa uses RoPE embeddings. </li> </ol> Stability. Problem of gradient vanishing or exploding as we pass through the layers of an rnn or transformer. Layer normalization (Ba 2016) is the traditional way to deal with this issue. The intuition is that it normalizes the outputs of each attention layer to a consistent range, preventing too much variance in the scale of outputs. LayerNorm(X;g,b)=σxg⋅(X−μx)+b Here, X∈Rd×n is the output of an attention layer of embedding dimension d and sequence length n, μx and σx are the element-wise mean and standard deviation respectively across the time positions, and g and b are learnable vectors of dimension d. Hence, if X has very large or small values, the normalization process standardizes the range of values. The parameters g and b allow the model flexibility to shift the values to a different part of the space. A simplification of layer norm is RMSNorm (Zhang and Sennrich 2019). It removes the mean and bias terms but does not hurt performance empirically. It is used in Llama. The only learnable parameters per layer is g. RMS(X)=n1i=1∑nxi2<svg xmlns="http://www.w3.org/2000/svg" width='400em' height='3.1968em' viewBox='0 0 400000 3196' preserveAspectRatio='xMinYMin slice'><path d='M702 80H40000040 H742v3062l-4 4-4 4c-.667.7 -2 1.5-4 2.5s-4.167 1.833-6.5 2.5-5.5 1-9.5 1 h-12l-28-84c-16.667-52-96.667 -294.333-240-727l-212 -643 -85 170 c-4-3.333-8.333-7.667-13 -13l-13-13l77-155 77-156c66 199.333 139 419.667 219 661 l218 661zM702 80H400000v40H742z'/></svg>RMSNorm(X)=RMS(X)X⋅g Residual Connections. Add an additive connection between input (to an attention layer) and the output. Residual(X,f)=f(X)+X Where f is the attention layer function. It prevents vanishing gradients (since we get X) and also allows f to focus on learning the difference from the input. In the self-attention context, having the residual connection also prevents the need for tokens to attend to themselves, since that is provided by X itself. Post vs Pre Layer Norm (Xiong 2020). This paper found that applying LayerNorm after the attention layer (and the residual connection) broke the residual connection, because we have LayerNorm(f(X)+X), which means that we are no longer guaranteed to get the input X. This led to some instability in transformer training. Instead, they found it is better to apply f(LayerNorm(X))+X, which led to more stable training, since the residual connection is preserved. Activation functions. <ul> <li>Vaswani used ReLU(x)=max(0,x)</li> <li>LLaMA uses Swish or SiLU (Hendricks and Gimpel 2016), which is Swish(x;β)=x⋅σ(β⋅x) Looks a lot like ReLU but avoids the zero gradient problem.</li> </ul> Transformers are powerful but fickle - Vaswani 2017 used Adam with learning rate increase and decrease (warm up). This is no longer that necessary after the pre-layer norm (Xiong 2020). AdamW (Loshchilov and Hutter 2017) is now more popular. It applies weight decay for regularization to Adam and corrects the previous implementation which applied the regularization incompletely. AdamW is thus theoretically more stable than Adam. In summary, some comparisons. The Mamba paper found that the LLaMA architecture is 10x more efficient in terms of scaling law compared to the original Vaswani. <div class="table-wrapper"><table><thead><tr><th style="text-align: center"></th><th style="text-align: center">Vaswani</th><th style="text-align: center">LLaMA</th></tr></thead><tbody> <tr><td style="text-align: center">Norm Position</td><td style="text-align: center">Post</td><td style="text-align: center">Pre</td></tr> <tr><td style="text-align: center">Norm Type</td><td style="text-align: center">LayerNorm</td><td style="text-align: center">RMSNorm</td></tr> <tr><td style="text-align: center">Non-linearity</td><td style="text-align: center">ReLU</td><td style="text-align: center">SiLU</td></tr> <tr><td style="text-align: center">Position-Encoding</td><td style="text-align: center">Sinusoidal</td><td style="text-align: center">RoPE</td></tr> </tbody></table> </div> <h1 id="lecture-6-decoding-strategies"><a class="header" href="#lecture-6-decoding-strategies">Lecture 6: Decoding Strategies</a></h1> What is an LLM? It is a model that defines a conditional probability distribution of a sequence given some input sequence X. P(Y∣X)=j=1∏JP(yj∣X,y1,...,yj−1) The nice thing about the conditional distribution is that we get some notion of the model's confidence about the next token to generate. The problem with the conditional distribution is hallucination, since models generally assign some small but non-zero probability to incorrect tokens, even if all the pre-training data is factual. See Kalai and Kempala 2023. Ancestral Sampling is to sample the next token based on the indicated confidence by the model. The nice thing is that the resultant generations follow exactly the distribution of the model. yj∼P(yj∣X,y1,...,yj−1) The problem with ancestral sampling is the long tail problem. Most language models have around 30k tokens, and the probabilities from the long tail adds up, so that there's a somewhat good chance of sampling something really unlikely. The obvious solution to this problem is to ignore the long tail and only sampling from the top-k most probably tokens. This is top-k sampling. This results in only tokens that the model is somewhat confident in. Alternatively, we could only sampling from the top-p probability mass. This is called top-p or nucleus sampling. This is to account for the case where top-k sampling is not so desirable because if most of the probability mass is only on say 3 tokens, we may only want to sample from them. Another alternative is epsilon sampling, where we only sample tokens with some minimum probability. This ensures that we only sample from tokens where the model is somewhat confident. Another strategy is to modify the "peakiness" of the data by controlling the distribution temperature. This is done by modifying the scaling factor for the final softmax layer. This allows the user to put more weights on the top results (<code>temperature < 1.0</code>) for factual answers, or spread the weights out more by increasing the temperature for say story generaion. Contrastive Decoding is a newer idea - the idea is that we use a smaller model to improve the performance of a larger model. Instead of just decoding from the larger model's distribution, we contrastively decode by choosing outputs where the expert model thinks are more likely than the smaller model. The intuition is that both the big and small model are degenerate in similar ways (e.g. keep repeating itself), but the expert has knowledge which the smaller model does not have. Hence taking the probability difference log(pexpert)−log(pamateur) helps to eliminate the degenerate cases and produce better results. <h2 id="mode-seeking-decoding-methods"><a class="header" href="#mode-seeking-decoding-methods">Mode-Seeking Decoding Methods</a></h2> Instead of sampling, another approach is to try to maximize the probability of generation, i.e. mode-seeking. Y^=argmaxy P(Y∣X) Greedy decoding chooses the most likely token at each step. yj=y′argmaxP(y′∣X,y1,...,yj−1) However, greedy decoding does not guarantee finding the most likely sequence. e.g. <code>the</code> is often the most likely word, but choosing <code>the</code> would exclude many other sentences that may be more likely. Hence Beam Search, which ensures that we don't miss a high-probability sequence "hidden" behind a lower-probability prefix. This is a form of breadth-first search, where we maintain a few options at any point in the search. <ul> <li>First, we explore the top 3 next tokens at time step 1</li> <li>Next, we explore the top 3 next tokens at time step 2 from each branch, leading to 9 options</li> <li>Then, we prune down to the top 3 paths from the 9 options and repeat the process for time step 3</li> </ul> In practice, beam search often results in a very non-diverse set, i.e. sentences that are very similar to each other. Hence we may want to introduce diversity into the process. Diverse beam search modifies the scoring when pruning beams to avoid choosing overly similar beams. The similarity score can be as simple as word jaccard similarity between pairs. Stochastic beam search modifies the next token selection to sampling instead of using top greedy decodings. <h2 id="minimum-bayes-risk"><a class="header" href="#minimum-bayes-risk">Minimum Bayes Risk</a></h2> The question is: Do we actually want the generations with the highest probability? In general, we do find that outputs with low probability tend to be worse than those with high probability. However, when we compare amongst the top outputs where probabilities are quite close, it becomes less clear. Say e.g. we performed beam search and found the top 3 sequences: <pre><code>the cat sat down - 0.3 the cat ran away - 0.25 the cat sprinted off - 0.2 </code></pre> <code>the cat sat down</code> is the highest probability sequence, but the combined probability mass is higher for the idea that "the cat left the area". So the idea is that we might prefer generations that have high agreement with other sequences (high probability and low risk). y^=y′∈Yhargmaxy∈Ye∑G(y,y′) In the equation above, Ye refers to a random sample from the model, say <code>100 samples</code>. Yh refers to our hypothesis space, supposedly the <code>top 10 outputs</code>. The risk function G measures the similarity of each candidate y′ against the samples. For example, a risk function could be ROUGE score, which measures the n-gram overlap between two sequences. Generally, MBR is high performance but high cost - even if G is <code>ROUGE-1</code> (i.e. unigram overlap), it significantly outperforms beam search or greedy search. Other MBR variants: output ensembling. Post-ensemble (Kobayashi 2018) compares pairwise embedding similarity between outputs acoss models and chooses outputs with highest average similarity. self-consistency (Wang 2023) prompts for an answer using chain of thought, samples multiple outputs, and extracts the answer from each sample (ignoring the explanations). The most frequently generated answer is then chosen. <h2 id="constrained-generation"><a class="header" href="#constrained-generation">Constrained Generation</a></h2> Sometimes we want to impose some constraints on the outputs, e.g. we want the model to suggest some hobbies but we do not want <code>climbing</code>, or more properly, say we want to omit toxic texts. Options: <ul> <li> Ask the model to exclude <code>climbing</code>: often does not work </li> <li> Logit Manipulation. Set the logit for the token(s) corresponding to <code>climbing</code> to be 0. This often messes up because there may be many synonyms or ways to express the same thing and its impossible to enumerate them </li> <li> Sample and Discard. We set up a new discriminator model which predicts whether a sequence corresponds to the idea of <code>climbing</code> or not. This is an easier task and often people will initialize the model from the original language model and train it to predict this idea from a small set of fine-tuning data. We can then get the generative model to generate a few samples and keep only samples where the predicted probability of the idea to avoid is low. Another variant of this is FUDGE (Yang and Klein 2021), where we multiply the generative probability of the next token p(yj) by the discriminator probability that the new sequence will belong to the idea that we desire (e.g. formality). The chosen token will be that which maximizes the combined score. </li> <li> RLHF. The alignment fine-tuning using RLHF may be viewed as a way to do constrained generation. An interesting paper that discusses RLHF as bayesian inference is <a href="https://ar5iv.labs.arxiv.org/html/2205.11275">Korbak 2022</a>. Instead of fine-tuning, one way is to do reward-augmented decoding (Deng and Raffel 2023). The idea is to have a reward model that modifies the generative probabilities based on rewards. (?) </li> </ul> <h2 id="human-in-the-loop-decoding"><a class="header" href="#human-in-the-loop-decoding">Human In the Loop Decoding</a></h2> Some strategies to incorporate human intervention in the generation: <ul> <li>Interleaved text. Model generates some text, human continues, then back to the model etc.</li> <li>Fine-grained replacement. Human selects some part of the generated text, and twiddles some knobs (e.g. "more descriptive" etc.)</li> <li>Choosing outputs. Model generates a few options, and human chooses one.</li> </ul> We could also use a model in the loop. One idea is Tree of thought prompting, which is somewhat analogous to beam search. We have the model generate a few sentences at a time, with a few samples. An external model then judges and chooses the best branches to continue generation. <h2 id="practical-considerations"><a class="header" href="#practical-considerations">Practical Considerations</a></h2> To increase decoding speed, one method is speculative decoding. We generally generate tokens with a small model, but when the small model is very uncertain, the small model will generate topk samples and a large model will pick the next token. This can speed up generation significantly. There are many libraries for fast decoding, e.g. vLLM, Outlines, disco etc. General takeaway is that a lot can be done at decoding time without needing to fine-tune the original model. <h1 id="lecture-7-prompting-strategies"><a class="header" href="#lecture-7-prompting-strategies">Lecture 7: Prompting Strategies</a></h1> Basic Prompting. Append a textual string to be beginning of the sequence and let the model complete. Some models are trained as chatbots and require a specific prompting template, e.g. GPT, but in the backend its simply formatted into a string and passed to the model. The important thing is that we need to follow the template that the model was trained on, otherwise performance could suffer. Post-processing refers to formatting the returned output from the LLM. E.g. ChatGPT supports markdown rendering, so if you ask it to generate a table it will generate it and then render it as a markdown table. Another form of post-processing is output selection, i.e. we extract the part of the output that we are actually interested in. We could do so by extracting keywords (e.g. <code>fantastic</code>, <code>terrible</code>) for a sentiment analysis task. An interesting phenomenon for predicting labels is that getting the model to predict <code>positive</code> or <code>negative</code> vs 1-5 labels, the former will do better. The intuition behind this is to think about the data that the model was trained on. It is likely that the model has seen many more reviews with the works <code>positive</code> or <code>excellent</code> vs numeric labels, so it might do better with those types of labels. Few-shot prompting (Brown 2021) basically injects a few examples of the task together with the instruction. One thing to take note of is that LLMs (especially smaller ones) are sensitive to small changes in the in-context examples: <ul> <li>Example ordering (Lu 2021)</li> <li>Label balance (Zhang 2022)</li> <li>Label coverage (Zhang 2022)</li> </ul> Effects of few-shot prompting are also sometime counter-intuitive. For example, replacing correct labels with random labels sometimes barely hurts the accuracy of the task. This suggests that the few-shot prompts are more for getting the structure of the response correct rather than learning the desired labelling logic. Sometimes, more demonstrations can also hurt accuracy (this may be due to the longer context length confusing the model). Chain of thought prompting (Wei 2022) basically tries to get the model to explain its reasoning before making an answer. The original idea was to include reasoning steps in the few-shot prompts to get the model to do likewise, and it found that this significantly improved the accuracy of the model. One interpretation for why this works is that it provides the model with adaptive computation time to generate the correct answer. e.g. a simple question like <code>1+1=</code> may be answered immediately but some complex logical question might require several steps, and the reasoning step allows the model to generate whatever it needs to get the answer right. The next step is unsupervised chain of thought prompting (Kojima 2022), which basically found that we can get the same results by just appending <code>let's think step by step</code> to the prompt, without requiring the few-shot examples. Another idea is that structuring outputs as computer programs can help (Madaan 2022). e.g. if we want the model to output a Direct Acyclic Graph, we could represent the output as a python graph class object. The reason that this works is perhaps because programs are highly structured and there is a lot of code in the pre-training data. Another useful method is to get the model to output <code>JSON</code> format, which it has also seen a lot of. Another idea is to have program-aided language models (Gao 2022). This allows the LLM to call a code interpreter or calculator to compute answers. This works especially well for numeric questions. <h2 id="prompt-engineering"><a class="header" href="#prompt-engineering">Prompt Engineering</a></h2> One thing of take note of is that the format should match that of a trained model. e.g. leaving out the space after the colon <code>Passage:text</code> can lead to severe performance degradation. Changing the casing to <code>PASSAGE: text</code> can also degrade performance. We can also do automatic prompt generation. e.g. use another model to paraphrase our prompt and then select the best response out of the samples. Another approach is gradient-based, where we try out different prompt words and choose the best prompt words based on some kind of loss. These types of methods can result in highly non-human sequences that somehow produce the best results, but they can also be exploited to elicit harmful responses. Another method along these lines is Prefix tuning (Li and Liang 2021), where they train an embedding prefix that is appended to the transformer weights in each layer according to the task of interest. This is akin to LORA methods which train additional weights that are appendable to the model. One way to view prompting is to view it as a human-interpretable prior to the model, which can be easier than fine-tuning. <h1 id="lecture-8-fine-tuning-and-instruction-tuning"><a class="header" href="#lecture-8-fine-tuning-and-instruction-tuning">Lecture 8: Fine-tuning and Instruction Tuning</a></h1> The general framework is that language models are pre-trained on the semi-supervised language modelling task on a very large corpus, and then fine-tuned on a downstream task. There are two paradigms for doing this: <ul> <li> Multi-task learning. The standard multi-task learning framework is to train the multiple tasks (language modelling and downstream task) simultaneously, e.g. by alternating mini batches or combining losses together. In <a href="https://ar5iv.labs.arxiv.org/html/2109.07437">Dery 2021</a>, the paper argues that learning jointly on the language modelling and end-task produces better results than if we pre-trained and then fine tuned. The intuition is that the model will be learning representations that are useful for both tasks. Another paper from Anthropic also shows that incorporating safety training at the beginning out-performs pre-training first and then fine-tuning to incorporate safety. </li> <li> Pre-train then fine-tune. However, because it is so expensive to perform the language modelling, usually pre-train and fine-tune is the actual paradigm that we follow. </li> </ul> Full fine-tuning means to simply continue training the language model on the training data that we have. This is also called supervised fine-tuning. This can be prohibitively expensive. Rajbhandari 2019 showed that training a 65B parameter with 16-bit mixed precision without any optimizations requires around 1TB of GPU memory, which is clearly not feasible. The simplest solution is to scale-up horizontally across multiple-GPUs. DeepSpeed ZeRo (Rajbhandari 2019) is a popular framework for doing so: <ul> <li>Stage 1: partitioning the optimizer state does not hurt optimization speed much. This brings down memory per device from 120GB to 31GB across 12 devices.</li> <li>Stage 2: in addition, it partitions the gradients. This brings memory further down to 16GB.</li> <li>Stage 3: in addition, it partitions the parameters. This brings memory down to 1.9GB but severely impacts computation speed.</li> </ul> An alternative strategy is to only fine tune a smaller set of parameters. Adapters (Houlsby 2019) is one way to do this. The idea is to add an adapter block after each attention block. The adapter block down-projects the embedding dimensionality to something small like <code>16</code>, passes it through a non-linearity, then up-projects back to the original dimension. Each adapter block only uses <code>2 x model_dim x adapter_dim</code> parameters. There are generally two benefits to parameter-efficient fine-tuning methods: <ul> <li>They are much more memory efficient. This is because we only need to back-propagate gradients on nodes which are on the computation path between the adapter block sto the final loss function</li> <li>They are more robust to over-fitting on the small set of fine-tuning data</li> </ul> An extension to Adapters is Adapter Fusion (Pfeiffer 2020). The idea is that instead of a single adapter after each attention block, we have multiple adapters, each trained on a different task. We then add an AdapterFusion block, which is basically a multi-head attention over each adapter block, so that it can learn to choose which adapters to use automatically. LoRA (Hu 2021) is very similar conceptually to adapters, with the important difference that it does not have any non-linearity. The idea is that we express the fine-tuned weights as follows (reference: <a href="https://cameronrwolfe.substack.com/p/easily-train-a-specialized-llm-peft">Cameron Wolfe's blog post</a>): Wft=Wpt+ΔW The goal is to learn ΔW with a low rank adaptation, so that it is parameter-efficient. Suppose for simplicity that Wft∈Rd×d. We may approximate ΔW:=A⋅B, where A∈Rd×r,B∈Rr×d. We can then simply freeze Wpt and modify the forward pass to become Wpt+A⋅B, and fine-tune the parameters for A and B. r can be as small as <code>8</code> or <code>2</code>, leading to a very parameter-efficient fine-tuning method. A is initialized with small random values, whilst B is initialized as zero, ensuring that we begin the finetuning process with the model's original outputs. The efficiency of LoRA compared to full fine-tuning is significant. For <code>d=512</code>, <code>r=8</code>, the efficiency is around <code>3%</code>. d22dr=d2r The reason LoRA has caught on is two-fold: <ol> <li>It does not require modifying the original model structure - we simply modify the <code>state_dict</code> of the original model by adding the ΔW</li> <li>Adapters incur additional inference latency due to the additional layers. LoRA has no additional latency at all</li> </ol> <h1 id="lecture-8-reinforcement-learning-and-human-feedback"><a class="header" href="#lecture-8-reinforcement-learning-and-human-feedback">Lecture 8: Reinforcement Learning and Human Feedback</a></h1> <a href="https://www.youtube.com/watch?v=hCnKbwPntrE&list=PL8PYTP1V4I8D4BeyjwWczukWq9d8PNyZp&index=8">Fall 2024 - Reinforcement Learning and Human Feedback</a> So far we optimize model using maximum likelihood of next token. This paradigm has some problems: <ul> <li>Problem #1: Some mistakes are worse than others, and we want to penalize more egregious mistakes more</li> <li>Problem #2: The target labels in MLE can be bad (e.g. toxic internet stuff, disinformation)</li> <li>Problem #3: Exposure bias. At training time, the prefix sequence is guaranteed to be good / normal. At generation, the prefix sequence can potentially become strange. <ul> <li>One example is repeating words. In normal language, when a word is repeated twice, it is more likely that the word repeats again. Hence language models learn this behaviour.</li> </ul> </li> </ul> So how do we measure how "good" an output is? <ul> <li>Objective assessment. Have an annotated correct answer and match against it. This approach is good for math problems or problems with objective answers. <ul> <li>GSM8K (Cobbe 2021) is a famous math problem dataset</li> </ul> </li> <li>Human evaluation.</li> </ul> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="database-course"><a class="header" href="#database-course">Database Course</a></h1> Based on CMU 15-445/645 Intro to Database Systems course taught by Andy Pavlo and Jignesh Patel in Fall 2023. The <a href="https://www.youtube.com/playlist?list=PLSE8ODhjZXjbj8BMuIrRcacnQh20hmY9g">youtube playlist</a>. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="lecture-1"><a class="header" href="#lecture-1">Lecture 1</a></h1> Course about how to design and implement a database management system. Textbook: Database System Concepts by Silberschatz, Korth and Sudarshan. Agenda: <ul> <li>Database Systems Background</li> <li>Relational Model</li> <li>Relational Algebra</li> <li>Alternative Data Models</li> </ul> A database is an organized collection of inter-related data that models some aspect of the real world. Databases are the core component of most computer applications. e.g. an excel spreadsheet is a database. SQLite is the most popular database as it is used in every cellphone. <h2 id="flat-file-strawman"><a class="header" href="#flat-file-strawman">Flat File Strawman</a></h2> Store our database as a csv file that we manage. e.g. each line corresponds to an artist, year, country etc. Problems: <ul> <li>Super slow to find the artist of interest as we use a for-loop to find each</li> <li>Super slow to update or delete an artist</li> <li>Data types are not stored on the csv file, we need to know which is an integer etc.</li> <li>Concurrent writes to the file are not supported</li> </ul> A Database Management System (DBMS) is a software that allows applications to store and analyze information in a database. A general purpose DBMS supports the definition, creation, querying, update and administrations of databases in accordance with some data model. Usually first choice is postgres or sqlite. A data model is a collection of concepts for describing the data in a database. A schema is a description of a particular collection of data, using a given data model. Examples of data models: <ol> <li>Relational</li> <li>Key / Value</li> <li>Graph</li> <li>Document / XML / Object</li> <li>Wide-Column / Column-family</li> <li>Array / Matrix / Vectors</li> <li>Hierarchical</li> <li>Network</li> <li>Multi-Value</li> </ol> 1 is the most common. 2-4 are considered NoSQL models (a loose term). 7-9 are obsolete. <h2 id="relational-model"><a class="header" href="#relational-model">Relational Model</a></h2> Early database applications were difficult to write. Every time the database schema or layout changed, IBM would need to rewrite database programs. Ted Codd devised the relational model to address this problem. The relational model is an abstraction. The relational model defines a database abstraction based on relations to reduce maintenance overhead. Key tenets: <ul> <li>Store database in simple data structures (relations)</li> <li>Physical storage left up to the DBMS implementation</li> <li>Access data through high-level language, DBMS figures out the best execution model</li> </ul> A relation is an unordered set that contain the relationship of attributes that represent entities. An n-ary relation is equivalent to a table with n columns. A tuple is a set of attributes values (also known as its domain) in the relation. The special value <code>NULL</code> is a member of every domain. A relation's primary key uniquely identifies a single tuple. Some DBMSs automatically create an internal primary key if a table does not define one. Primary key is a constraint that the DBMS will enforce to ensure no duplicates exist. A foreign key specifies that an attribute from one relation maps to a tuple in another relation. E.g. If we have an artist table with the artist id, and an album table with an artist column, the artist column is a foreign key. We can impose constraints on the database that must hold for any tuple. DBMS will then prevent any modification that could violate those constraints. Unique and foreign key constraints are the most common. e.g. <code>CREATE ASSERTION</code> in SQL. <h2 id="data-manipulation-languages-dml"><a class="header" href="#data-manipulation-languages-dml">Data Manipulation Languages (DML)</a></h2> There are two broad methods to store and retrieve information from a database: <ol> <li>Procedural: the query specifies the high level strategy to find the desired result based on sets. This uses relational algebra.</li> <li>Non-Procedural (Declarative): The query specifies only what data is wanted and not how to find it. This uses relational calculus.</li> </ol> <h2 id="relational-algebra"><a class="header" href="#relational-algebra">Relational Algebra</a></h2> Fundamental operations to retrieve and manipulate tuples in a relation. Based on set algebra (unordered lists with no duplicates). Each operator takes one or more relations as its inputs and outputs a new relation. We can thus chain operators together to create more complex operations. The operations are: <ul> <li>SELECT. Choose a subset of the tuples from a relation that satisfies a selection predicate (filter). Predicates act as filters to retain only tuples that fulfill the qualifying requirement. We can combine multiple predicates using conjunctions / disjunctions. <ul> <li>Syntax: σpredicate(R)</li> <li><code>SELECT * from TABLE where id="a"</code></li> </ul> </li> <li>PROJECTION. Generate a relation with tuples that contains only the specified attributes. E.g. re-arrange ordering, manipulate values (+,− etc.) and remove unwanted attributes. <ul> <li>Syntax: πA1,A2,...(R)</li> <li>Example: <code>SELECT b_{id} - 100</code></li> </ul> </li> <li>UNION. Generate a relation that contains all tuples that appear in one or both input relations. Note that R and S must have the same schema. <ul> <li>Syntax: (R⋃S)</li> <li>Example: <code>(SELECT * from R) UNION (SELECT * from S)</code></li> </ul> </li> <li>INTERSECTION. Generate a relation that contains only the tuples that appear in both of the input relations. <ul> <li>Syntax: (R⋂S)</li> <li>Example: <code>(SELECT * from R) INTERSECT (SELECT * from S)</code></li> </ul> </li> <li>DIFFERENCE. Generate a relation that contains only the tuples that appear in the first and not the second of the input relations. <ul> <li>Syntax: (R−S)</li> <li>Example: <code>(SELECT * from R) EXCEPT (SELECT * from S)</code></li> </ul> </li> <li>JOIN. Generate a relation that contains all tuples that are a combination of two tuples (one from each input relation) with a common value for one or more attributes. <ul> <li>Syntax: (R ∞ S)</li> <li>Example: <code>SELECT * FROM R NATURAL JOIN S;</code></li> </ul> </li> <li>Rename ρ</li> <li>Assignment R←S</li> <li>Duplicate Elimination δ</li> <li>Aggregation γ</li> <li>Sorting τ</li> <li>Division (R÷S)</li> </ul> Relational algebra defines an ordering of the high level steps of how to compute a query. E.g. <ul> <li>σbid=102(R ∞ S) vs R ∞ (σbid=102(S)). The former will do a huge join before filtering, whereas the latter filters first before joining, which is much better.</li> </ul> Instead of specifying the exact operations, DBMS allow us to state the high level answer we want, and the DBMS decides on the exact operations and the ordering to perform. This abstracts away the need for the user to know which operations are more efficient. Note that the relational model is independent of the query language implementation, although SQL is the de-facto standard (with many variants). <h2 id="document-data-model"><a class="header" href="#document-data-model">Document Data Model</a></h2> A collections of record documents containing a hierarchy of named field / value pairs. A field's value can be a scalar type, array, or another document. Modern implementations use <code>JSON</code>. Main reason for this model is to avoid relational object impedance msimatch, i.e. relational databases store data in rows with relationships between tables, but in object oriented languages like Python data is stored in objects with nested attributes, which could result in inefficient queries when we try to map between the two. In contrast, Document Databases store data in a nested json which closely resembles the object-oriented approach, making it easier to work with. The down side is that we could end up storing a lot of duplicate data in the json objects. <ul> <li>Examples: MongoDB, RavenDB, DynamoDB etc.</li> </ul> <h2 id="vector-data-model"><a class="header" href="#vector-data-model">Vector Data Model</a></h2> One-dimensional arrays used for nearest neighbour search, used for semantic search on embeddings generated by transformer models. Native integration with modern ML tools etc. At their core, these are just systems that use specialized indexes (e.g. Meta FAISS) to perform NN search quickly. <ul> <li>Examples: Pinecone, Weaviate, Milvus etc.</li> </ul> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="lecture-2-modern-sql"><a class="header" href="#lecture-2-modern-sql">Lecture 2: Modern SQL</a></h1> 1971, the first relational query language called SQUARE was created. 1972 SEQUEL (Structured English Query Language) was created. SQL was added to ANSI standard in 1986. Current standard is SQL:2023. <ul> <li>SQL:2023 - property graph queries, multi-dim arrays</li> <li>SQL:2016 - JSON, polymorphic tables etc.</li> </ul> Relational languages: <ul> <li>Data Manipulation Langauge (DML)</li> <li>Data Definition Language (DDL)</li> <li>Data Control Language (DCL)</li> </ul> Important: <SQL is based on bags> (with duplicates) not sets (no duplicates). We should try to do everything on the database, in one big query. Example database: <ul> <li>student table: <code>sid</code>, <code>login</code>, <code>gpa</code> etc.</li> <li>enrolled table: <code>sid</code>, <code>cid</code> (course id), <code>grade</code></li> <li>course table: <code>cid</code>, <code>name</code></li> </ul> Aggregates: <ul> <li><code>AVG(col)</code>. e.g. <code>SELECT AVG(s.gpa) FROM student as s</code></li> <li><code>MIN(col)</code></li> <li><code>MAX(col)</code></li> <li><code>COUNT(col)</code>. e.g. <code>SELECT COUNT(LOGIN) as cnt FROM student WHERE login LIKE '%@cs</code>. Equivalently, <code>COUNT(1)</code>. Count number of rows where their login matches the pattern.</li> </ul> Groupby: Get average gpa by course id. <pre><code class="language-sql">SELECT AVG(s.gpa), e.cid FROM enrolled as e JOIN student AS s ON e.sid = s.sid GROUP BY e.cid </code></pre> String operations. <ul> <li><code>LIKE</code> is used for string matching. <code>%</code> matches any substring, including empty strings. <code>_</code> matches any one character.</li> <li>SQL-92 defines string functions.</li> </ul> Window functions. Perform a sliding calculation across a set of tuples that are related. Like an aggregation but tuples are not grouped into a single output tuple. <pre><code class="language-sql">SELECT ... FUNC-NAME(...) OVER (...) FROM TABLE_NAME </code></pre> e.g. Get row number per course id. <pre><code class="language-sql">SELECT cid, sid ROW_NUMBER() OVER (PARTITION BY cid) FROM enrolled ORDER BY cid </code></pre> Nested queries. Invoke a query inside of another query to compose more complex computations. These are often difficult for the DBMS to optimize. e.g. This one below is a join written as a nested query. <pre><code class="language-sql">outer query -> SELECT name FROM student WHERE sid IN (SELECT sid FROM enrolled) <- inner query </code></pre> e.g. Get the names of students in '15-445': <pre><code class="language-sql">SELECT name FROM student WHERE sid IN ( SELECT sid FROM enrolled WHERE cid = '15-445' ) </code></pre> Lateral joins. <code>LATERAL</code> operator allows us to reference tables preceding the expression. e.g. below, the second expression can reference t1. <pre><code class="language-sql">SELECT * FROM (SELECT 1 AS X) AS t1, LATERAL (SELECT t1.x+1 AS y) AS t2; </code></pre> Common Table Expressions. Provides a way to write auxiliary statements for use in a larger query, i.e. a temp table just for one query. <pre><code class="language-sql">WITH cteName (col1, col2) as ( SELECT 1, 2 ) SELECT col1 + col2 FROM cteName </code></pre> Demonstration of CTEs: Find student record with the highest id that is enrolled in at least one course. We use the maxId in the temporary table below. <pre><code class="language-sql">WITH cteCourse (maxId) AS ( SELECT MAX(sid) FROM enrolled ) SELECT name FROM student, cteSource WHERE student.sid = cteSource.maxId </code></pre> We can also use recursion with CTE. Print the sequence of numbers from 1 to 10. <pre><code class="language-sql">WITH RECURSIVE cteSource (counter) AS ( (SELECT 1) UNION (SELECT counter + 1 FROM cteSource WHERE counter < 10) ) SELECT * FROM cteSource </code></pre> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="lecture-3-database-storage-1"><a class="header" href="#lecture-3-database-storage-1">Lecture 3: Database Storage 1</a></h1> Class will shift gears towards how to implement a DBMS. Focus is on a single node machine first, then move to concurrency. Today's focus is on the <Disk Manager>. Focus on <disk-based architecture>, i.e. the DBMS assumes that the primary storage location of the database is on non-volatile disk. The DBMS's components manage the movement of data between volatile (e.g. RAM) and non-volatile storage. Storage hierarchy. Higher rank is faster, smaller, expensive. <ol> <li>CPU Registers</li> <li>CPU Caches</li> <li>DRAM (i.e. memory)</li> <li>SSD</li> <li>HDD</li> <li>Network Storage</li> </ol> Items 1-3 are volatile, i.e. random access, byte-addressable (we can fetch specific bytes). Items 4-6 are non-volatile, i.e. sequential access, block-addressable. For non-volatile storage, we need to fetch data one block at a time, and accessing data in one contiguous block sequentially is much faster than trying to access blocks scattered around. CPU refers to items 1 and 2, Memory refers to item 3, Disk refers to items 4-6. Some new storage category: Fast network storage is in between HDD and SSD. Persistent Memory is between SSD and DRAM. DBMSs need to exploit the sequential access on non-volatile memory to maximize speed of retrieval. Hence how we store data contiguously matters in the design of the DBMS. Access times: <pre><code>- 1 ns L1 Cache Ref 1 sec - 4 ns L2 Cache Ref 4 sec - 100 ns DRAM 100 sec - 16,000 ns SSD 4.4 hours - 2,000,000 ns HDD 3.3 weeks - 50,000,000 ns Network Storage 1.5 years </code></pre> <h2 id="system-design-goals"><a class="header" href="#system-design-goals">System Design Goals</a></h2> The overall goal is to make it appear like we have more memory than we do: <ol> <li>We want to allow the DBMS to manage databases that exceed the amount of memory available.</li> <li>Reading/ writing to disk is expensive, so it must be managed carefully to avoid large stalls.</li> <li>Random access on disk is much slower, so DBMS wants to maximize sequential access.</li> </ol> <h2 id="disk-oriented-dbms"><a class="header" href="#disk-oriented-dbms">Disk oriented DBMS</a></h2> On disk, we have database files stored in pages. Each page is a block of data. When an execution engine makes a query, it loads a page into buffer pool in memory. The buffer pool returns a 64-bit pointer to the memory location where the page is located. The execution engine interprets the layout of the page, updates it. The new page is then written back into the database file on disk. <h2 id="why-not-use-the-os"><a class="header" href="#why-not-use-the-os">Why not use the OS?</a></h2> The DBMS could potentially use memory mapping <code>mmap</code> to store the contents of a file (or page) in the address space of a program. The benefit is that we can tap on the OS to decide which pages to load into physical memory and when to swap pages out. The downside is that the lack of control over memory management can lead to stalling and bad performance. How the OS works: <ol> <li>Suppose we have Page 1, ..., Page 4 on disk. There's only enough space in physical memory to load 2 pages.</li> <li>The OS represents this to the program as virtual memory, where Page 1, ..., Page 4 are available</li> <li>When the program touches Page 2, the OS will load Page 2 into physical memory and return a pointer to the program</li> <li>Suppose the program now touches Page 3 and Page 4</li> <li>The OS needs to decide which page to evict from physical memory</li> </ol> So the problems with using <code>mmap</code> to handle I/O for the DBMS: <ol> <li>Transaction safety. The OS can flush dirty pages to disk at any time, even mid-write. We can get corrupted data on disk.</li> <li>I/O stalls. The DBMS does not know which pages are in memory, so the OS could stall due to page faults (fetching a page not in memory).</li> <li>Error handling. Difficult to validate pages. Any access can cause a <code>SIGBUS</code> that the DBMS needs to handle.</li> <li>Performance issues. The OS has its own scheduling and data structures, and can contend with the DBMS's priorities.</li> </ol> There are some companies that use <code>mmap</code>: <ul> <li>RavenDB</li> <li>ElasticSearch</li> <li>QuestDB</li> <li>MongoDB (moved away from <code>mmap</code>)</li> </ul> So DBMS almost always wants to control things itself: <ul> <li>Flushing dity pages to disk in the correct order</li> <li>Specialized pre-fetching</li> <li>Buffer replacement policy</li> <li>Thread / process scheduling</li> </ul> Paper on this: <a href="https://db.cs.cmu.edu/papers/2022/cidr2022-p13-crotty.pdf">Are you sure you want to Use MMAP in your DBMS?</a> For database storage, we need to deal with two problems: <ol> <li>How the DBMS represents the database in files on disk</li> <li>How the DBMS manages its memory and moves data back and forth from disk</li> </ol> <h2 id="file-storage"><a class="header" href="#file-storage">File Storage</a></h2> The DBMS stores a database as one or more files on disk typically in a proprietary format. The OS doesn't know anything about the contents of these files. There are portable file formats (e.g. parquet etc.). Early systems in 1980s further used custom filesystems on raw block storage. Most newer DBMSs do not do this. The <storage manager> is responsible for maintaining a database's files. It organizes the files as a collection of pages, and tracks the data read / written to pages, and tracks the available space in each page. A DBMS does not typically maintain multiple copies of a page on disk. This happens above or below the storage manager. A <page> is a fixed-size block of data. It can contain tuples, meta-data, indexes, logs etc. Most systems do not mix page types. Some systems require a page to be self-contained, i.e. all the information needed to understand the data in the page is contained within the page itself (e.g. Oracle). Each page is given a unique identifier. The DBMS uses an indirection layer to map page IDs to physical locations (e.g. in memory or in S3 or whatever). There are 3 different notions of what a page is: <ol> <li>Hardware page (typically <code>4KB</code>). A hardware page is the largest block of data that the storage device can guarantee failsafe writes (i.e. if it says it wrote <code>4KB</code>, it actually happened).</li> <li>OS Page (usually <code>4KB</code>, but can be much larger)</li> <li>Database Page (<code>512B</code> to <code>32KB</code>)</li> </ol> Why might a larger page size be a good idea? Sequential data access. If we request a page of <code>16KB</code>, it is one call to the OS and data is stored sequentially. If we request 4 pages of <code>4KB</code> each, it is potentially random access. But the downside is that writing to a larger page is slower than writing to a small one. <h2 id="page-storage-architecture"><a class="header" href="#page-storage-architecture">Page Storage Architecture</a></h2> How do DBMSs map page IDs to physical location? Different approaches: <ul> <li>Heap file organization <- most common</li> <li>Tree file organization</li> <li>Sequential / Sorted File Organization (ISAM)</li> <li>Hashing file organization</li> </ul> A heap file is an unordered collection of pages with tuples that are stored in random order. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="rl-course"><a class="header" href="#rl-course">RL Course</a></h1> Lecture notes based on <a href="https://www.youtube.com/watch?v=2pWv7GOvuf0&list=PLqYmG7hTraZDM-OYHWgPebj2MfCFzFObQ">David Silver's Reinforcement Learning course</a>. Textbooks: <ul> <li>An Introduction to Reinforcement Learning - Sutton and Barto</li> <li>Algorithms for Reinforcement Learning, Szepesvari</li> </ul> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="lecture-1-introduction-1"><a class="header" href="#lecture-1-introduction-1">Lecture 1: Introduction</a></h1> <a href="https://youtu.be/2pWv7GOvuf0?feature=shared">Lecture 1</a> Reinforcement Learning is at the centre of many disciplines such as engineering, psychology, economics, neuroscience because it deals with the science of decision making. E.g. in neuroscience, one major part is the dopamine neurotransmitter which resembles RL. Mathematics - operations research. Economics - bounded rationality. What makes RL different from other ML paradigms? <ul> <li>There is no supervisor, only a reward signal. There is no clear correct action to take, only a reward</li> <li>The feedback is delayed, not instantaneous</li> <li>Time really matters in RL (non iid data)</li> <li>Agent's actions affect the subsequent data it receives</li> </ul> Examples: <ul> <li>Fly stunt manoeuvres in a helicopter</li> <li>Play backgammon</li> <li>Manage an investment portfolio</li> <li>Make a humanoid robot walk</li> </ul> Rewards. A reward Rt is a scalar feedback signal. Indicates how well agent is doing at step t. The agent's job is to maximize cumulative reward. RL is based on the reward hypothesis. <blockquote> Reward Hypothesis. All goals can be described by the maximization of expected cumulative reward. </blockquote> Is a scalar feedback sufficient? David's argument is that in order to pick one action, we must always be able to compare and rank two actions. Hence we must have a way of ordering actions and it boils down to a scalar score. Examples of rewards: <ul> <li>Fly stunt manoeuvres in a helicopter: <ul> <li>+ve reward for following desired trajectory</li> <li>-ve reward for crashing</li> </ul> </li> <li>Backgammon: <ul> <li>+ve / -ve reward for winning / losing a game</li> </ul> </li> <li>Manage investment portfolio <ul> <li>+ve reward for each $ in bank</li> </ul> </li> </ul> Sequential decision making. Goal: select actions to maximize total future reward. We have to plan ahead, because actions may have long term consequences. It may be better to sacrifice immediate reward to gain more long term reward. Greedy approach does not work in RL. Formalism. At each step t, the agent: <ul> <li>Executes action At</li> <li>Receives observation Ot from environment</li> <li>Receives reward Rt from environment</li> </ul> The environment: <ul> <li>Receives action At</li> <li>Emits observation Ot</li> <li>Emits reward Rt</li> </ul> The history is the sequence of observations, actions, rewards. Ht=A1,O1,R1,...,At,Ot,Rt What happens next depends on the history: <ul> <li>The agent selects the next action based on the history</li> <li>The environment selects observation and reward to emit</li> </ul> This is a very generic framework or formalism that can handle all types of real world scenarios. State is the information used to determine what happens next. We do not want to have to load the entire history (e.g. stream of video frames) to make a decision. Formally, state is a function of the history: St=f(Ht) The simplest example of state is to e.g. just take the observation at the last timestamp (this worked for Atari games). The environment state Ste is the environment's private representation. What state the environment is in. That is, the data that the environment uses to pick the next observation / reward. The environment state is not usually visible to the agent or the algorithm. The agent state Sta is the agent's internal representation, i.e. it is the information used by reinforcement learning algorithms. It can be any function we choose of the history: Sta=f(Ht). A more mathematical definition of state. An information state (or Markov state) contains all useful information from the history. <blockquote> Definition. A state S_t is Markov if and only if P(St+1∣St)=P(St+1∣S1,...,St) </blockquote> The idea may be stated as "the future is independent of the past given the present". Once the state is known, the history may be thrown away. i.e. The state is a sufficient statistic of the future. In the helicopter example, the markov state might be something like position, velocity, angle, angular velocity etc. Having known all this information, it does not matter where the helicopter was 10 minutes ago, as we already have all the information we need to make an optimal decision next. Two trivial statements to show that there always exists a markov state: <ul> <li>The environment state Ste is Markov by definition</li> <li>The history Ht is Markov, albeit not a useful one</li> </ul> Fully observable environments. Full observability means that agent directly observes the environment state. The nice case. Ot=Sta=Ste Formally, this is a markov decision process (MDP). Partially observable environments. Partial observability means that agent indirectly observes the environment. e.g. <ul> <li>Robot with camera doesn't know its absolute location</li> <li>A trading agent only observes current prices</li> <li>A poker playing agent only observes public cards, not hidden cards</li> </ul> Now, agent state is not the same as the environment state. Formally this is a partially observable markov decision process. So now we need the agent to construct its own state representation Sta, e.g.: <ul> <li>Remember the complete history Sta=Ht</li> <li>Build beliefs of the environment state, i.e. Sta=(P(ste)=s1,..P(Ste)=sn). So we have a probability distribution over possible states that we believe the environment is in.</li> <li>Recurrent neural network, i.e. Sta=σ(St−1aWs+OtWo). Use a linear transformation to combine the current observation with previous time step to get current time step.</li> </ul> <h2 id="inside-an-rl-agent"><a class="header" href="#inside-an-rl-agent">Inside an RL agent</a></h2> An RL agent may include these components: <ul> <li>Policy: the agent's behaviour function</li> <li>Value function: how good is each state or action</li> <li>Model: the agent's representation of the environment</li> </ul> Policy. A policy is the agent's behaviour. It is a map from state to action, e.g. <ul> <li>Deterministic policy, a=π(s)</li> <li>Stochastic policy, π(a∣s)=P(A=a∣S=s)</li> </ul> Value function is a prediction of future reward, or the expected future reward. It is used to evaluate the goodness or badness of states. vπ(s)=Eπ[RT+γRt+1+γ2Rt+2+...∣St=s] Note that the value function is subscripted by π. This indicates that the agent's expected future reward depends on its current policy. If the policy is causing a robot to fall a lot, the expected future reward is probably low. Model A model predicts what the environment will do next. There are usually 2 parts to the model: <ul> <li>Transitions P predicts the next state (dynamics of the environment)</li> <li>Rewards R predicts the next immediate reward</li> </ul> Formally, Pss′a=P(S′=s′∣S=s,A=a)Rsa=E[R∣S=s,A=a] Note that having a model is optional, there are many model-free methods in RL. Maze example. We have a grid maze. The rewards are −1 per time step. The actions are N, E, S, W. The states are the agent's current location. <ul> <li>An example of a deterministic policy would be to have a fixed arrow direction for the action to take in any given grid that the agent is in.</li> <li>An example of value function is to have a number in each position showing the number of steps it would take to get to the end (but negated as we lose value each time step)</li> </ul> A taxonomy of RL agents: <ul> <li>Value based - only using a value function, the policy is implicit</li> <li>Policy based - just have a policy function, no value function</li> <li>Actor critic - have both a value function and policy function, and try to get best of both worlds</li> </ul> Another categorization is model free vs model approach: <ul> <li>Model free means we just go straight to policy and/or value function <ul> <li>No model</li> </ul> </li> <li>Model based approach. Try to model the environment and world first, then build the policy accordingly</li> </ul> <h2 id="learning-and-planning"><a class="header" href="#learning-and-planning">Learning and Planning</a></h2> Two fundamental problems in sequential decision making: <ul> <li>Reinforcement learning <ul> <li>The environment is initially unknown</li> <li>The agent interacts with the environment</li> <li>The agent improves its policy</li> </ul> </li> <li>Planning <ul> <li>A model of the environment is known and provided to the agent</li> <li>The agent performs internal computations with the model without any external interaction</li> <li>e.g. if when playing an atari game, we gave the model access to an atari emulator. The model then knows if it takes action a, the emulator will be in this state etc. The model can then plan ahead, build a search tree etc.</li> </ul> </li> </ul> <h2 id="exploration-and-exploitation"><a class="header" href="#exploration-and-exploitation">Exploration and Exploitation</a></h2> Reinforcement learning is like trial and error learning. But it is possible to miss out on exploring steps that can lead to more reward. Exploration means choosing to gives up some known reward in order to find out more information about the environment. Exploitation exploits known information to maximize reward. Examples: <ul> <li>Restaurant selection: <ul> <li>Exploitation: go to your favourite restaurant</li> <li>Exploration: try a new restaurant</li> </ul> </li> <li>Online banner advertisements <ul> <li>Exploitation: show the most successful advert</li> <li>Exploration: show a different and new advert</li> </ul> </li> </ul> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="lecture-2-markov-decision-processes"><a class="header" href="#lecture-2-markov-decision-processes">Lecture 2: Markov Decision Processes</a></h1> <a href="https://www.youtube.com/watch?v=lfHX2hHRMVQ&list=PLqYmG7hTraZDM-OYHWgPebj2MfCFzFObQ&index=2">Lecture 2: Markov Decision Processes</a> Markov Decision Processes formally describe an environment for Reinforcement Learning. <ul> <li>The environment is fully observable, i.e. the current state fully characterizes the process</li> <li>Almost all RL problems can be characterized as an MDP</li> <li>Even continuous things like Optimal Control</li> <li>Partially observable cases can be formulated as MDPs</li> <li>Bandits are MDPs with one state</li> </ul> The Markov Property is central to MDPs. "The future is independent of the past given the present." State Transition Matrix. For a Markov state s and successor state s′, the state transition probability is defined as: Pss′=P(St+1=s′∣St=s) The state transition matrix P defines transition probabilities from all states s to all successor states s′. P=<svg xmlns="http://www.w3.org/2000/svg" width='0.667em' height='4.200em' viewBox='0 0 667 4200'><path d='M403 1759 V84 H666 V0 H319 V1759 v600 v1759 h347 v-84 H403z M403 1759 V0 H319 V1759 v600 v1759 h84z'/></svg>P11⋮P11⋯⋱⋯P1n⋮Pnn<svg xmlns="http://www.w3.org/2000/svg" width='0.667em' height='4.200em' viewBox='0 0 667 4200'><path d='M347 1759 V0 H0 V84 H263 V1759 v600 v1759 H0 v84 H347z M347 1759 V0 H263 V1759 v600 v1759 h84z'/></svg> Each row of the transition matrix sums to 1. Markov Process. A Markov Process is a memoryless random process, i.e. a sequence of of random states S1,S2,... with the Markov Property. <blockquote> Definition. A Markov Process (or Markov Chain) is a tuple <S,P>, where: (i) S is a (finite) set of states (ii) P is a state transition probability matrix (iii) Pss′=P(St+1=s′∣St=s) </blockquote> Example of a Markov Process. A student can transit from Class 1 to Class 2 to Class 3, Pass or Sleep or Pub based on transition probabilities. We can sample episodes for the markov chain. E.g. one episode may be <code>C1 C2 C3 Pass Sleep</code>. The transition probability matrix may look something like the below. Note that <code>Sleep</code> is the terminal state, so its self-probability is <code>1.0</code>. P=<svg xmlns="http://www.w3.org/2000/svg" width='0.667em' height='9.600em' viewBox='0 0 667 9600'><path d='M403 1759 V84 H666 V0 H319 V1759 v6000 v1759 h347 v-84 H403z M403 1759 V0 H319 V1759 v6000 v1759 h84z'/></svg>C1C2C3PassPubFBSleepC100000.20.10C20.50000.400C300.8000.400Pass000.60000Pub000.40000FB0.500000.90Sleep00.201.0001<svg xmlns="http://www.w3.org/2000/svg" width='0.667em' height='9.600em' viewBox='0 0 667 9600'><path d='M347 1759 V0 H0 V84 H263 V1759 v6000 v1759 H0 v84 H347z M347 1759 V0 H263 V1759 v6000 v1759 h84z'/></svg> A markov reward process is a markov chain with values. <blockquote> Definition. A Markov reward process is a tuple <S,P,R,γ>: <ul> <li>S is a finite set of states</li> <li>P is a transition probability matrix, Pss′=P[St+1=s′∣St=s]</li> <li>R is a reward function, Rs=E[Rt+1∣St=s]</li> <li>γ is a discount factor, γ∈[0,1]</li> </ul> </blockquote> <blockquote> Definition. The return Gt is the total discounted reward from time step t. Gt=Rt+1+γRt+2+...=∑k=0∞γkRt+k+1 </blockquote> There is no expectation because Gt is one sample run of the Markov reward process. We'll take expectation later to get the expected return over infinite runs. <ul> <li>Note that the discount factor is the present value of future rewards. 0 implies maximally short-sighted and 1 implies maximally far-sighted.</li> <li>The value of receiving reward R after k+1 time steps is γkR</li> <li>This setup values immediate reward above future reward</li> </ul> Most Markov reward and decision processes are discounted. Why? <ul> <li>We do not have a perfect model, so the expected future rewards are more uncertain. Hence we put higher weights on immediate rewards.</li> <li>Avoids inifite returns in cyclic Markov processes</li> <li>If the reward is financial immediate rewards earn more interest than delayed rewards</li> <li>Animal / human behaviour shows preference for immediate rewards</li> <li>It is sometimes possible to use undiscounted Markov reward processes, e.g. if all sequences terminate</li> </ul> The value function v(s) gives the long-term value of state s. <blockquote> Definition. The state value function v(s) of an MRP is the expected return stating from state s: v(s)=E[Gt∣St=s] </blockquote> How do we compute the state value function? One way is to sample returns from the MRP. e.g. stating from S1=C1 and γ=1/2: <ul> <li><code>C1 C2 C3 Pass Sleep</code>: <code>-2.25</code></li> <li><code>C1 FB FB C1 C2 Sleep</code>: <code>-3.125</code></li> </ul> Consider if we set γ=0. Then the value function v(s)=Rs, i.e. the value is just the immediate reward. Now the important Bellman equation for MRPs: v(s)=E[Rt+1+γv(St+1)∣St=s] It essentially tells us that the value function can be decomposed into two parts: <ul> <li>Immediate reward Rt+1</li> <li>Discounted value of successor state γv(St+1)</li> </ul> v(s)=E[Gt∣St=s]=E[Rt+1+γRt+2+γ2Rt+3+...∣St=s]=E[Rt+1+γ(Rt+2+γRt+3+...)∣St=s]=E[Rt+1+γGt+1∣St=s]=E[Rt+1+γv(St+1)∣St=s]=Rs+γs′∑Pss′v(s′) <ul> <li>Note that in the second-to-last line, the argument inside v(St+1) is a random variable, to express the fact that the state at time t+1 is random.</li> <li>Note that both Gt+1 and v(St+1) are random variables, which express the value function at each possible state at time step t+1.</li> <li>Gt becomes v(St+1) due to the law of iterated expectations. Recall that E[X]=E[E[X∣Y]]. (Not very sure exactly how this works out.)</li> </ul> To dig into bellman equation a bit more. Use a 1-step look ahead search. We start at state s, we look ahead one step and integrate over the probabilities of the next time step. Hence we get v(s)=Rs+γ∑s′Pss′v(s′). We can use the bellman equation to verify if our value function is correct. Taking the value at a particular state, we can check if it is indeed the sum of the immediate reward and the weighted sum of values in all possible next steps. The Bellman equation can be expressed concisely using matrices, v=R+γPv <svg xmlns="http://www.w3.org/2000/svg" width='0.667em' height='4.200em' viewBox='0 0 667 4200'><path d='M403 1759 V84 H666 V0 H319 V1759 v600 v1759 h347 v-84 H403z M403 1759 V0 H319 V1759 v600 v1759 h84z'/></svg>v(1)⋮v(n)<svg xmlns="http://www.w3.org/2000/svg" width='0.667em' height='4.200em' viewBox='0 0 667 4200'><path d='M347 1759 V0 H0 V84 H263 V1759 v600 v1759 H0 v84 H347z M347 1759 V0 H263 V1759 v600 v1759 h84z'/></svg>=<svg xmlns="http://www.w3.org/2000/svg" width='0.667em' height='4.200em' viewBox='0 0 667 4200'><path d='M403 1759 V84 H666 V0 H319 V1759 v600 v1759 h347 v-84 H403z M403 1759 V0 H319 V1759 v600 v1759 h84z'/></svg>R1⋮Rn<svg xmlns="http://www.w3.org/2000/svg" width='0.667em' height='4.200em' viewBox='0 0 667 4200'><path d='M347 1759 V0 H0 V84 H263 V1759 v600 v1759 H0 v84 H347z M347 1759 V0 H263 V1759 v600 v1759 h84z'/></svg>+γ<svg xmlns="http://www.w3.org/2000/svg" width='0.667em' height='4.200em' viewBox='0 0 667 4200'><path d='M403 1759 V84 H666 V0 H319 V1759 v600 v1759 h347 v-84 H403z M403 1759 V0 H319 V1759 v600 v1759 h84z'/></svg>P11⋮Pn1⋯⋱⋯P1n⋮Pnn<svg xmlns="http://www.w3.org/2000/svg" width='0.667em' height='4.200em' viewBox='0 0 667 4200'><path d='M347 1759 V0 H0 V84 H263 V1759 v600 v1759 H0 v84 H347z M347 1759 V0 H263 V1759 v600 v1759 h84z'/></svg><svg xmlns="http://www.w3.org/2000/svg" width='0.667em' height='4.200em' viewBox='0 0 667 4200'><path d='M403 1759 V84 H666 V0 H319 V1759 v600 v1759 h347 v-84 H403z M403 1759 V0 H319 V1759 v600 v1759 h84z'/></svg>v(1)⋮v(n)<svg xmlns="http://www.w3.org/2000/svg" width='0.667em' height='4.200em' viewBox='0 0 667 4200'><path d='M347 1759 V0 H0 V84 H263 V1759 v600 v1759 H0 v84 H347z M347 1759 V0 H263 V1759 v600 v1759 h84z'/></svg> The bellman equation is a linear equation and can be solved directly using matrix inversion. v(1−γP)vv=R+γPv=R=(I−γP)−1R The complexity due to the matrix inversion if O(n3) for n states, which is not feasible for a large number of states. There are many iterative methods which are more efficient: <ul> <li>Dynamic programming</li> <li>Monte Carlo evaluation</li> <li>Temporal Difference learning</li> </ul> <h2 id="markov-decision-process"><a class="header" href="#markov-decision-process">Markov Decision Process</a></h2> So far it has been a building block. The MDP is what we really use. A Markov Decision Process (MDP) is a markov reward process with decisions (actions). It is an environment in which all states are Markov. <blockquote> Definition. A Markov Decision Process is a tuple <S,A,P,R,γ>: <ul> <li>S is a finite set of states</li> <li>A is a finite set of actions</li> <li>P is a transition probability matrix with Pss′a=P[St+1=s′∣St=s,At=a]</li> <li>R is a reward function, Rsa=E[Rt+1∣St=s,At=a]</li> <li>γ is a discount factor</li> </ul> </blockquote> Note that the transition probabilities and reward functions now also depend on an action, in which we can have some agency now. We can choose actions to influence the reward and values. <blockquote> Definition. A policy π is a distribution over actions given states, π(a∣s)=P[At=a∣St=s] </blockquote> A policy fully defines the behaviour of an agent. Some properties of a policy: <ul> <li>It only depends on the current state (not the history)</li> <li>The policy does not depend on the time step t (i.e. stationary)</li> </ul> We can still obtain the optimal policy because of the markov property - the current state captures all relevant information to make the optimal decision. <blockquote> Definition. The state-value function vπ(s) of an MDP is the expected return starting from state s, and following policy π: vπ(s)=Eπ[Gt∣St=s] </blockquote> <blockquote> Definition. The action-value function qπ(s,a) of an MDP is the expected return starting from state s, taking action a, and following policy π: qπ(s,a)=Eπ[Gt∣St=s,At=a] </blockquote> Bellman Expectation Equation. The state-value function can again be decomposed into immediate reward plus discounted value of the successor state. vπ(s)=Eπ[Rt+1+γvπ(St+1)∣St=s] Similarly we can do so for the action-value function, by inserting the chosen action: qπ(s,a)=Eπ[Rt+1+γqπ(St+1,At+1)∣St=s,At=a] From a given state, we have a value function attached to that state, i.e. vπ(s). From this state, we have some possible actions to take. The policy determines the probability distribution over which action to take. With each action comes an action-value function qπ(s,a). Hence we have: vπ(s)=a∈A∑π(a∣s)qπ(s,a) Another way to look at it. We start with having chosen a particular action. Having chose a particular action, the environment will determine the particular state I end up in (based on the transition probability matrix P). Hence we have: qπ(s,a)=Rs+γs′∈S∑Pss′avπ(s′) Now we can stitch these two perspectives together. Starting from a particular state, we can write vπ(s) in terms of qπ, then write qπ in terms of vπ again. This will allow us to get a recursive relationship of vπ(s) in terms of vπ(s′) and allow us to solve the equation. The bellman expectation equation for vπ(s) is thus: vπ(s)=a∈A∑π(a∣s)(Rsa+γs′∈S∑Pss′avπ(s′)) The math is expressing a simple idea: that the value at a particular state s is the weighted sum of values from all possible actions we take under the current policy π. The value of each action is in turn affected by the reward function and the transition probability that determines the state we end up in after taking a particular action. Similarly, we can do the same by starting at an action instead of a state. The bellman expectation equation for qπ is thus: qπ(s,a)=Rsa+γs′∈S∑Pss′aa′∈A∑π(a′∣s′)qπ(s′,a′) <h2 id="optimal-value-function"><a class="header" href="#optimal-value-function">Optimal Value Function</a></h2> So far we have been defining the dynamic process of the MDP, but have not tried solving the optimization problem. We will turn to this now. <blockquote> Definition. The optimal state-value function v∗(s) is the maximum value function over all policies: v∗(s)=πmaxvπ(s) The optimal action-value function q∗(s,a) is the maximum action-value function over all policies: q∗(s,a)=πmaxqπ(s,a) </blockquote> The MDP problem is solved once we find q∗. We thus need some algorithms to systematically find q∗. Define a partial ordering over policies: π≥π′ if vπ(s)≥vπ′(s)∀s <blockquote> Theorem. For any MArkov Decision Process: <ul> <li>There exists an optimal policy π∗ that is better than or equal to all other policies, i.e. π∗≥π,∀π</li> <li>All optimal policies achieve the optimal value function, i.e. vπ∗(s)=v∗(s),∀s</li> <li>All optimal policies achieve the optimal action-value function, i.e. qπ∗(s,a)=q∗(s,a),∀s,a</li> </ul> </blockquote> How do we find the optimal policy? An optimal policy can be found trivially by maximizing over q∗(s,a), if we knew it. That is, we always pick the action a with the highest q(s,a) value. Hence if we have q∗, we have π∗. π∗(a∣s)={10if a=arga′∈Amaxq∗(s,a′)otherwise Intuitively, we find the optimal policy by starting at the end (resting), and iteratively look backward. This is the same kind of intuition for the Bellman optimality equations. The optimal value of being in a state s is the highest value action we can take in that state. Note that we use q∗ instead of a generic q because we are choosing from the optimal action-value function. v∗(s)=amaxq∗(s,a) The optimal value of an action a is the weighted sum of values of states that we can end up in after taking the action. Note that in this step, we do not get to choose an action - the transition probabilities will determine what state we end up in after taking actions a: q∗(s,a)=Rsa+s′∈S∑Pss′av∗(s′) Finally, stitching these two equations together, we get the bellman optimality equation for v∗: v∗(s)=amax[Rsa+s′∈S∑Pss′av∗(s′)] How do we solve the bellman optimality equations? It is now non-linear due to the <code>max</code> function, so we cannot solve it with matrix inversion as before. There is no closed form solution in general, but there are many iterative solution methods: <ul> <li>Value iteration</li> <li>Policy iteration</li> <li>Q-learning</li> <li>Sarsa</li> </ul> <blockquote> Intuition. The core idea behind the bellman equations is to break down a complex sequential decision problem into a series of simpler, recursive steps. Imagine we are at a particular point in time and in a particular state. The bellman equations tell us that if we can assume that we will act optimally for all future steps after this action, then the problem of finding the best current action becomes trivial - we simply choose the action that yields the highest expected value (based on assuming future optimality). To actually start unravelling the equations and solving them, we start from the termination point of a process (where the assumption of future optimality trivially holds) and work backwards. </blockquote> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="lecture-3-planning-by-dynamic-programming"><a class="header" href="#lecture-3-planning-by-dynamic-programming">Lecture 3: Planning by Dynamic Programming</a></h1> <a href="https://www.youtube.com/watch?v=Nd1-UUMVfz4&list=PLqYmG7hTraZDM-OYHWgPebj2MfCFzFObQ&index=3">Lecture 3: Planning by Dynamic Programming</a> What is dynamic programming? <ul> <li>Dynamic: sequential or temporal component to the problem</li> <li>Programming: optimising a mathematical "program", i.e. a policy</li> </ul> It is a method for solving complex problems, by breaking them into subproblems that are simpler to solve. Dynamic Programming is a very general solution method for problems which have two properties: <ul> <li>Optimal substructure: the pricniple of optimality applies. The optimal solution can be decomposed into subproblems. <ul> <li>e.g. to find the shortest path from A to B, we can find the shortest path from A to midpoint, and then find the shortest path from midpoint to B, and then combine the paths together.</li> </ul> </li> <li>Overlapping subproblems. The subproblems need to recur many times and solutions can be re-used. <ul> <li>e.g. if we have the shortest path from midpoint to B, we can reuse that to find the shortest path from C to B if it traverses the midpoint as well.</li> </ul> </li> </ul> Note that Markov Decision Processes satisfy both properties: <ul> <li>The bellman equations decompose the large problem into recursive steps</li> <li>The value functions for a particular state cache the sub-solutions and are re-used</li> </ul> Planning by dynamic programming. Planning is a different problem from RL. Someone tells us the dynamics of the MDP, and we try to solve it. <ul> <li>Assume full knowledge of the MDP</li> <li>Used for planning in an MDP</li> <li>We can use this for prediction: <ul> <li>e.g. input: an MDP <S,A,P,R,γ> and policy π</li> <li>The output of this planning step is to output the value function vπ</li> </ul> </li> <li>We can also use this for control: <ul> <li>Input: MDP <S,A,P,R,γ></li> <li>Output: optimal value function v∗, i.e. we want to find the best policy π∗</li> </ul> </li> </ul> <h2 id="policy-evaluation"><a class="header" href="#policy-evaluation">Policy Evaluation</a></h2> Problem: we want to evaluate a given policy π to see how good it is. The solution is to iteratively apply the bellman expectation. <ul> <li>Let v1 be a vector representing the value at all states</li> <li>We use the bellman equation to update v1→v2→...→vπ, i.e. it will converge to the true value function for this policy</li> <li>Do this using synchronous backups: <ul> <li>At each iteration step k+1</li> <li>For all states s∈S</li> <li>Update vk+1(s) using vk(s′), i.e. we use the estimate of the value function in the previous iteration to form the new estimate of the value function</li> <li>s′ is a successor state of s, i.e. the next states we can reach using an action from s</li> </ul> </li> <li>It can be proven that this algorithm will converge to vπ</li> </ul> How exactly do we update vk+1(s)? We use the bellman expectation equation from before. Intuitively, it is a one-step look ahead from the current state s to compute the value for s. vk+1(s)=a∈A∑π(a∣s)(Rsa+γs′∈S∑Pss′avk(s′)) Or in vector form: vk+1=Rπ+γPπvk Small Gridworld example. Suppose we have a 4x4 grid, in which the top-left and bottom-right grids are terminal states. The reward for any state is −1, and we can walk NSEW from any spot. If we walk off the grid, the action just does nothing. Now suppose we take a random walk (0.25 probability of walking any direction) and set γ=1. How would the value update algorithm look like? <ul> <li>We initialize all grids with value 0</li> <li>Recall that the update algorithm essentially adds the immediate reward and a discounted sum of the value function from the previous iteration</li> <li>At k=1, every spot will be updated to −1 because the reward is −1, and v0(s)=0 ∀s. Except the two terminal states which remain at value 0 by definition</li> </ul> If we continue to update like that, it will converge to the true value function for π. Also note that with this value function vπ, if we take a greedy approach to devise a new policy, we can obtain the optimal policy. <blockquote> Idea: A lousy policy can be used to devise a better policy after computing the value function. </blockquote> <h2 id="policy-iteration"><a class="header" href="#policy-iteration">Policy Iteration</a></h2> We did policy evaluation in the previous section, i.e. finding the true value function vπ for a given policy. Now in this section we want to optimize and find the best policy. Two step process: <ul> <li>Policy Evaluation. Given a policy π, we first evaluate the policy π, finding vπ(s)=E[Rt+1+γRt+2+...∣St=s]</li> <li>Policy improvement. We take a greedy approach and choose the best action at each state: π′=greedy(vπ)</li> </ul> Typically, we need many rounds of iteration of this process to converge. But the process of policy iteration always converges to π∗. Specifically: <ul> <li>v converges to v∗</li> <li>π converges to π∗</li> </ul> Somewhat realistic toy example from Sutton and Barto: <ul> <li>Suppose we have two locations, maximum of 20 cars at each location</li> <li>Each day, we get to move up to <code>5</code> cars between locations overnight</li> <li>Reward: <code>$10</code> for each car rented (can only rent if we have enough cars)</li> <li>Transitions: every day, a random number of cars are requested and returned at each location (governed by a poisson distribution) <ul> <li>Location A: request∼Poisson(3), return∼Poisson(3)</li> <li>Location B: request∼Poisson(4), return∼Poisson(2)</li> </ul> </li> </ul> Naturally we expect the optimal policy to involve moving cars from location A to location B. Using the policy iteration process, we get convergence to the optimal policy in 4 steps. Note that since this is a planning problem, we do know the underlying probability mechanisms, which allows us to compute the value function. Now we can show that this policy iteration process converges to the optimal policy. <blockquote> Theorem. Policy iteration process converges to the optimal policy. Proof. First, consider a deterministic policy a=π(s). Now, we consider what happens if we change the policy by acting greedily wrt the value function of this policy, i.e.: π′(s)=a∈Aargmaxqπ(s,a) We see that taking the greedy action can only improve the policy, as expressed: qπ(s,π′(s))=a∈Amaxqπ(s,a)≥qπ(s,π(s))=vπ(s) Notes on the above statement: <ul> <li>On the 2nd line equality: Recall that qπ(s,a) is the value at state s if we took action a at s and then followed policy π thereafter. So it follows that qπ(s,π(s))=vπ(s), since we are just following policy π in the current step + future steps</li> <li>The statement is quite simply saying that π′(s)≥π(s), which leads to an improvement in qπ. Hence choosing the highest value action will improve the action value function (quite trivial).</li> </ul> Now we want to go from this somewhat trivial statement to show that the value function itself must improve with every step of policy iteration (not trivial at all!). The idea is to use a telescoping argument to show that this improves the value function, or vπ′(s)≥vπ(s): vπ(s)≤qπ(s,π′(s))=Eπ′[Rt+1+γvπ(St+1)∣St=s]≤Eπ′[Rt+1+γqπ(St+1,π′(St+1))∣St=s]≤Eπ′[Rt+1+γRt+2+γ2qπ(St+2,π′(St+2))∣St=s]≤Eπ′[Rt+1+γRt+2+...∣St=s]=vπ′(s) Some notes on the above: <ul> <li>We start with the trivial inequality expressed above</li> <li>The expression Eπ′ means taking expectation over possible trajectories under the policy where we take π′(s) in the current step, then follow policy π for the rest of the trajectory</li> <li>In line 2, we unpack qπ according to the Bellman equation, which simply splits up the q value into (i) the immediate reward and (ii) the expected value of our new state (expressed as a random variable S): <ul> <li>Note that Rt+1,Rt+2,... are rewards from taking the greedy action π′(s) at each step</li> <li>Note that vπ(St+n) is the random variable expressing the value we have at the next time step, but evaluated under the previous policy π. It is the previous π intead of π′ because we have access to the cached vπ up to this point.</li> </ul> </li> <li>In line 3, we apply the trivial inequality again to show that taking the greedy step at the next state will again improve the value</li> <li>In line 4, we again use the Bellman equation to unpck qπ</li> <li>We keep repeating the two steps until termination. What we have at the end is simply the value function of our new policy π′</li> </ul> We have shown that policy iteration must improve the value function with each iteration. Now what happens when improvements stop? We now have: qπ(s,π′(s))=vπ(s) Since qπ(s,π′(s))=maxa∈Aqπ(s,a), we have: vπ(s)=a∈Amaxqπ(s,a) This is simply the Bellman optimality equation. Satisfying the Bellman optimality equation means that we are in the optimal state (will show later). Hence we have shown that we get vπ(s)=v∗(s). </blockquote> Now, an observation is that policy iteration is quite wasteful. This is because we need to get the value function to converge fully to vπ before we take the greedy step to improve the policy. In most cases, this is unnecessary because the greedy policy would already improve even with an imperfect value function. Some ways to early stop policy evaluation to speed up this process: <ul> <li>Introduce a stopping condition once the value function does not change by much (ϵ-convergence of value function)</li> <li>Stop policy evaluation after k iterations <ul> <li>In the extreme case, if we stop policy evaluation after k=1 iterations, it is called value iteration</li> </ul> </li> </ul> <h2 id="value-iteration"><a class="header" href="#value-iteration">Value Iteration</a></h2> Moving into value iteration, but recall the fundamentals of dynamic programming. Observe that any optimal policy can be subdivided into two components: <ul> <li>An optimal first action A∗</li> <li>Followed by an optimal policy from successor state S′</li> </ul> <blockquote> Theorem. Principle of optimality. A policy π(a∣s) achieves the optimal value from state s, i.e. vπ(s)=v∗(s), if and only if for any state s′ reachable from s, π achives the optimal value from state s′, i.e. vπ(s′)=v∗(s′) </blockquote> This theorem seems a bit of a truism, but it will be used to build the idea of value iteration. Let us think of the value function as "caching" the solutions to subproblems. Now suppose we start "at the end" and assume we know the solution to all the subproblems v∗(s′) where s′ is all the states reachable from our current state s. Then we can solve immediately for v∗(s) by doing a one-step lookahead to all these states s′: v∗(s)←a∈AmaxRsa+γs′∈S∑Pss′av∗(s′) The statement above shows us how we can propagate the optimal value function from some states s′ to a new state s. So we can propagate the optimal value function across to all states as we continue to iterate. The way to think about it (using the small gridworld as example) is that the termination point (trivially) starts off with the optimal value function. After one step of update, the states next to the termination point will now have the optimal value function, and then the states next to these, and so on until we propagate through all states. Note that in contrast to policy iteration, where in the policy evaluation step we update the value function across all states based on the bellman expectation equation, in value iteration, we are updating the value in each state by choosing the optimal action. This is a key difference in how the two algorithms differ. The value iteration algorithm may be thought of as combining the (i) policy evaluation step and the (ii) greedy policy step from value itaration into one single step. So we have seen that value iteration iteratively applies the bellman optimality equation to update v1→v2→...→v∗. (Note: Useful to compare this update statement with the bellman expectation equation used for policy evaluation above). Convergence to v∗ will be proved later. The update equation is: vk+1(s)=a∈Amax(Rsa+γs′∈S∑Pss′avk(s′)) Note that another difference between value iteration and policy iteration is that there is no explicit policy in value iteration. Since we are only doing one step of policy evaluation and then immediately taking the greedy step, the value function we have may not correspond to any real policy. But this does not stop the algorithm from converging. The following table sums up the relationship between what we've learnt. <div class="table-wrapper"><table><thead><tr><th>Problem</th><th>Bellman Equation</th><th>Algorithm</th></tr></thead><tbody> <tr><td>Prediction</td><td>Bellman Expectation Equation</td><td>Iterative Policy Evaluation</td></tr> <tr><td>Control</td><td>Bellman Expectation Equation + Greedy Policy Improvement</td><td>Policy Iteration</td></tr> <tr><td>Control</td><td>Bellman Optimality Euqation</td><td>Value Iteration</td></tr> </tbody></table> </div> Some notes: <ul> <li>These algorithms are based on the state value function vπ(s)</li> <li>The complexity for m actions and n states is O(mn2) per iteration</li> <li>We could also apply the same algorithm to the action-value function qπ(s,a)</li> <li>But the complexity worsens to O(m2n2) per iteration</li> </ul> <h2 id="extensions-to-dynamic-programming"><a class="header" href="#extensions-to-dynamic-programming">Extensions to Dynamic Programming</a></h2> <ul> <li>The DP methods described so far used synchronous backups, i.e. we backup all states in parallel</li> <li>Asynchronous backs up state individually in any order, without updating all states in one step</li> <li>This can significantly reduce computation</li> <li>There are nice properties to show that it is guaranteed to converge if we still select all states in the way we update</li> </ul> Now, three simple ideas: <ul> <li>In place Dynamic Programming</li> <li>Prioritized Sweeping</li> <li>Real time dynamic programming</li> </ul> In place dynamic programming. A simple idea where we update the value function in-place rather than store it in a separate array. Original (store updates in a separate array): <ul> <li>For all s∈S: <ul> <li>vnew(s)←maxa∈A(Rsa+γ∑s′∈SPss′avold(s′))</li> </ul> </li> <li>vold←vnew</li> </ul> New (in place updates right away): <ul> <li>For all s∈S: <ul> <li>v(s)←maxa∈A(Rsa+γ∑s′∈SPss′av(s′))</li> </ul> </li> </ul> The method with in place updates has more recent updates to v, and thus often are a lot more efficient in convergence. Prioritized Sweeping. Since we are doing immediate updates to v, it begs the question: in what order should we update states? <ul> <li>One method is to use the magnitude of the bellman error to guide state selection: error=<svg xmlns="http://www.w3.org/2000/svg" width='0.333em' height='3.000em' viewBox='0 0 333 3000'><path d='M145 15 v585 v1800 v585 c2.667,10,9.667,15,21,15 c10,0,16.667,-5,20,-15 v-585 v-1800 v-585 c-2.667,-10,-9.667,-15,-21,-15 c-10,0,-16.667,5,-20,15z M188 15 H145 v585 v1800 v585 h43z'/></svg>a∈Amax(Rsa+γs′∈S∑Pss′av(s′))−v(s)<svg xmlns="http://www.w3.org/2000/svg" width='0.333em' height='3.000em' viewBox='0 0 333 3000'><path d='M145 15 v585 v1800 v585 c2.667,10,9.667,15,21,15 c10,0,16.667,-5,20,-15 v-585 v-1800 v-585 c-2.667,-10,-9.667,-15,-21,-15 c-10,0,-16.667,5,-20,15z M188 15 H145 v585 v1800 v585 h43z'/></svg></li> <li>The idea is that states with the largest bellman error are those states whose value functions will change the most, which will significantly change the dynamics of the system, so we should update them first</li> <li>This can be implemented efficiently by maintaining a priority queue</li> </ul> Real time dynamic programming. Idea is to select states that the real world agent is visiting to update. <ul> <li>Basically update the states that the agent is visiting right now</li> <li>After each time step we have St,At,Rt+1</li> <li>So we update the state St</li> </ul> Important note: DP uses full-width backups, whether we are doing sync or async updates. This means that we consider the max over every successor state and action. Also, we need full knowledge of the MDP dynamics to compute. For large problem spaces, one single backup may be too expensive to compute. Hence in subsequent lectures we will consider sample backups. Can use the contraction mapping theorem to prove convergence etc. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="model-free-prediction"><a class="header" href="#model-free-prediction">Model Free Prediction</a></h1> <a href="https://www.youtube.com/watch?v=PnHCvfgC_ZA&list=PLqYmG7hTraZDM-OYHWgPebj2MfCFzFObQ&index=4">Lecture 4: Model Free Prediction</a> Methods in which no one tells us the environment, as opposed to before. <h2 id="introduction-1"><a class="header" href="#introduction-1">Introduction</a></h2> <ul> <li>Lecture 3: planning by dynamic programming, in which we solve a known MDP</li> <li>Lecture 4: Model-free prediction, in which we estimate the value function of an unknown MDP <ul> <li>Unknown in the sense that we do not have access to the environment, only the interactions that the agent has with the environment</li> </ul> </li> <li>Lecture 5: Model-free control (or optimization), in which we optimize the value function of an unknown MDP</li> </ul> <h2 id="monte-carlo-reinforcement-learning"><a class="header" href="#monte-carlo-reinforcement-learning">Monte Carlo Reinforcement Learning</a></h2> MC methods learn directly from episodes of experience. <ul> <li>It is model-free in that there is no knowledge of MDP transitions / rewards.</li> <li>It only learns from complete episodes</li> <li>MC uses the simplest idea: the value function of a state is the average return from that state over many many runs</li> <li>Downside: MC can only be applied to episodic MDPs <ul> <li>Episodic meaning that all episodes must terminate and we get a return value</li> </ul> </li> </ul> The goal of MC Policy Evaluation is to learn vπ from episodes of experience under policy π: <ul> <li>An episode: S1,A1,R2,...,Sk∼π</li> </ul> Recall that: <ul> <li>The return is the total discounted reward: Gt=Rt+1+γRt+2+...+γT−1RT</li> <li>And the value function is the expected return given that we start at a given state: vπ(s)=Eπ[Gt∣St=s]</li> </ul> So the whole idea of MC policy evaluation is to replace the expected return function with the empirical mean return from observing many many episodes. There are two main methods of performing this: <ul> <li>First visit MC policy evaluation</li> <li>Every visit MC policy evaluation</li> </ul> Method 1: First visit MC policy evaluation. Algorithm for evaluating a given state s is: <ul> <li>At the first time step t where state s is visited in a given episode: <ul> <li>Increment counter N(s)←N(s)+1. N(s) is the count of episodes where s was visited.</li> <li>Increment total return S(s)←S(s)+Gt</li> <li>Value is estimated by mean return V(s)=S(s)/N(s)</li> </ul> </li> <li>By the law of large numbers, V(s)→vπ(s) as N(s)→∞</li> </ul> Note that Gt above is the total discounted reward from time step t onwards. Method 2: Every visit MC policy evaluation. The algorithm is identical to first visit, with the only difference in that we perform the increment step every time we visit state s. <h3 id="blackjack-example"><a class="header" href="#blackjack-example">BlackJack example</a></h3> <ul> <li>There are <code>200</code> unique states: <ul> <li>Current sum (<code>12</code> to <code>21</code>). If it's <code>11</code> or below, the action is automatically to <code>twist</code>.</li> <li>Dealer's showing card <code>ace to 10</code></li> <li>Do I have a "useable" ace?</li> </ul> </li> <li>Actions: <ul> <li><code>stick</code>: stop receiving cards</li> <li><code>twist</code>: take another card.</li> </ul> </li> <li>Reward for <code>stick</code>: <ul> <li><code>+1</code> if <code>our sum > dealer's sum</code></li> <li><code>0</code> if <code>our sum = dealer's sum</code></li> <li><code>-1</code> if <code>our sum < dealer's sum</code></li> </ul> </li> <li>Reward for <code>twist</code>: <ul> <li><code>-1</code> if <code>our sum > 21</code></li> <li><code>0</code> otherwise</li> </ul> </li> </ul> We can use MC policy evaluation algorithm to play <code>10,000</code> episodes of blackjack and compute the value function of each state, under a given policy. For e.g. a naive policy is to <code>stick</code> if <code>our sum >= 20</code>, otherwise <code>twist</code>. <h3 id="incremental-mean"><a class="header" href="#incremental-mean">Incremental Mean</a></h3> The mean of a sequence of values can be computed in an incremental algorithm: μk=k1j=1∑kxj=k1(xk+j=1∑k−1xj)=k1(xk+(k−1)μk−1)=μk−1+k1(xk−μk−1) The last line shows that at each step, we just need to adjust the running mean μ by a small quantity, which is the difference between the new observed value xk and the current mean μk−1. This is analogous to a gradient update. So applying this incremental mean algorithm to monte carlo updates. Recall that the value function V(s) is the mean return over episodes. Hence we can change the above MC algorithm to an incremental mean update: <ul> <li>After observing a given episode S1,A1,R2,...,ST: <ul> <li>For each state St with return Gt: N(St)V(St)←N(St)+1←V(St)+N(St)1(Gt−V(St))</li> <li>We may even replace the running count N(St) with a fixed step size α. This is the usual approach in non-stationary problem. This algorithm allows us to avoid keeping track of old episodes and just keep updating V(St). V(St)←V(St)+α(Gt−V(St))</li> </ul> </li> </ul> <h2 id="temporal-difference-learning"><a class="header" href="#temporal-difference-learning">Temporal Difference Learning</a></h2> TD methods are different from MC methods, in that we do not need to wait for full episodes to learn. <ul> <li>TD methods, like MC methods, learn directly from episodes of experience</li> <li>TD methods, like MC methods, are also model-free</li> <li>TD learns from incomplete episodes using bootstrapping</li> <li>TD updates a guess towards a guess</li> </ul> Goal remains the same: learn vπ online from experience under policy π. Simplest temporal difference learning algorithm: TD(0). <ul> <li>Update value V(St) toward estimated return Rt+1+γV(St+1) V(St)←V(St)+α(Rt+1+γV(St+1)−V(St))</li> <li>Rt+1+γV(St+1) is called the TD target -Rt+1+γV(St+1)−V(St) is called the TD error</li> </ul> Contrast this with incremental every-visit Monte Carlo which we saw earlier: <ul> <li>Update value V(St) toward actual return Gt V(St)←V(St)+α(Gt−V(St))</li> </ul> <h3 id="car-driving-example"><a class="header" href="#car-driving-example">Car Driving Example</a></h3> An analogy for understanding the difference between MC and TD methods. Imagine we are driving the car home. At the start, we expect to take <code>30 mins</code> for the journey. And then it rains, so we update our prediction to <code>40 mins</code>. And so on. Eventually, the final journey takes <code>43 mins</code>. <ul> <li>For MC method, we need to wait until the journey is complete, and then we update the value of our policy to <code>43 mins</code> each step of the way</li> <li>For TD method, we can immediately update the value function to the next prediction each step of the way</li> </ul> <h3 id="bias-variance-trade-off"><a class="header" href="#bias-variance-trade-off">Bias Variance Trade-off</a></h3> There is a bias variance trade-off between choosing MC or TD method for policy evaluation. <ul> <li>The return Gt=Rt+1+γRt+2+..+γT−1RT is an unbiased estimate of vπ(St)</li> <li>The oracle TD target Rt+1+γvπ(St+1) is also an unbiased estimate of vπ(St) <ul> <li>We know this from the bellman expectation equation</li> <li>But it requires access to the oracle vπ(St+1) which we do not have</li> </ul> </li> <li>The TD target Rt+1+γV(St+1) is a biased estimate of vπ(St) <ul> <li>This is because V is our current estimate of the value function which can be wildly wrong</li> </ul> </li> <li>Observe that the TD target has much lower variance than the return: <ul> <li>The return depends on many random actions, transitions, rewards through the entire run of the episode</li> <li>The TD target only depends on one random action, transition and reward <ul> <li>The value function V is a deterministic function</li> </ul> </li> </ul> </li> </ul> So to summarize: <ul> <li>MC has high variance and zero bias <ul> <li>So it has good convergence properties, even with function approximation later</li> <li>It is not very sensitive to the initial value</li> <li>Very simple to understand and use</li> </ul> </li> <li>TD has low variance but some bias <ul> <li>Usually it is much more efficient than MC</li> <li>TD(0) can be proven to converge to vπ(s) using a table lookup</li> <li>But with function approximation convergence is not always guaranteed</li> <li>More sensitive to the initial value</li> </ul> </li> </ul> What is function approximation? This will be covered later on. But in general, we have been looking at v(s) as a table lookup for each state. This is not feasible for problems with large state spaces, hence we need to learn a function to approximate v(s) for all states. MC vs TD empirical example <ul> <li>TD generally converges faster than MC</li> <li>But if the step size α is too larger, TD may not fully converge as it will oscillate</li> </ul> So far we have seen that both MC and TD converge as the number of episodes goes to infinity. <ul> <li>That is, V(s)→vπ(s) as episodes→∞</li> <li>But what if we only have a limited number of K episodes to learn from?</li> <li>For example, what if we are repeatedly sampling episode k∈[1,K]?</li> </ul> <h3 id="ab-example"><a class="header" href="#ab-example">AB Example</a></h3> A simple example to illustrate difference between MC and TD in the finite data case. Suppose we have 6 episodes: <ul> <li><code>A, 0, B, 0</code></li> <li><code>B, 1</code></li> <li><code>B, 1</code></li> <li><code>B, 1</code></li> <li><code>B, 1</code></li> <li><code>B, 0</code></li> </ul> What would we think V(A),V(B) are? <ul> <li>If we use MC, then V(A)=0,V(B)=4/6. V(A) is 0 because we only encountered one episode involving state A and the reward was 0.</li> <li>If we use TD, then V(A)=4/6,V(B)=4/6. V(A) is 4/6 because we observed a <code>100%</code> probability of transiting from A→B, so the value of A (without discounting) is the same as value of B due to bootstrapping.</li> </ul> In more precise terms: <ul> <li>MC converges to the solution which minimizes the mean squared error <ul> <li>i.e. minimizes the divergence from observed returns mink=1∑Kt=1∑Tk[gtk−V(stk)]2</li> <li>In the AB example above, this sets V(A)=0</li> </ul> </li> <li>TD converges to the solution of the maximum likelihood Markov model <ul> <li>i.e. it converges to the MDP that best fits the data</li> <li>In the AB example, V(A)=4/6</li> </ul> </li> </ul> In summary: <ul> <li>TD exploits the markov property, so it is usually more efficient in markov environments, where we can rely on states to encode information</li> <li>MC does not exploit the markov property, so it is usually more efficient in non-markov environments, e.g. partial observability etc.</li> </ul> So far we have looked at 3 types of backup: <ul> <li>Monte Carlo Backup: we sample one entire trajectory / episode from the agent's interactions with the environment till termination V(St)←V(St)+α(GT−V(St))</li> <li>Temporal Difference Backup: we sample one step lookahead and then update parameters V(St)←V(St)+α(Rt+1+γV(St+1)−V(St))</li> <li>Dynamic Programming Backup: we look ahead one step, but because we have access to the environment, we can compute the expectation over all possible next steps. V(St)←Eπ[Rt+1+γV(St+1)]</li> </ul> This gives us two dimensions to categorize our algorithms: <ul> <li>Bootstrapping: the update involves an estimate (e.g. our value function) <ul> <li>MC does not bootstrap</li> <li>DP bootstraps</li> <li>TD bootstraps</li> </ul> </li> <li>Sampling: we use sampling instead of a full-width expectation / search <ul> <li>MC samples</li> <li>DP does not sample</li> <li>TD samples</li> </ul> </li> </ul> <h2 id="td-lambda"><a class="header" href="#td-lambda">TD Lambda</a></h2> TD Lambda is a generalization of the above trade-off. We let TD target look n steps into the future before updating. If we look forward ∞ number of steps, it becomes monte carlo learning. Specifically, for n=1,2,∞, our returns are: <ul> <li>n=1: Gt(1)=Rt+1+γV(St+1)</li> <li>n=2: Gt(2)=Rt+1+γRt+2+γ2V(St+2)</li> <li>n=∞: Gt(∞)=Rt+1+γRt+2+...+γT−1RT. We can see this corresponds to MC update, without use of value function V at all</li> </ul> So the n-step return is: Gt(n)=Rt+1+γRt+2+...+γn−1Rt+n+γnV(St+n) And the n-step TD learning update is: V(St)←V(St)+α(Gt(n)−V(St)) What is the best n? It is a highly sensitive parameters that depends on the problem, α etc. Hence a proposal is made to average the returns from each time step, up to step n. For example, we could average 21Gt(2)+21Gt(4). This averaging would make the algorithm much more robust to step size n. The common way to perform a weighted average of returns is to use exponential λ decay, such that returns with a longer look-ahead window are weighted less. This algorithm is called TD-λ. Specifically: Gtλ=(1−λ)n=1∑∞λn−1Gt(n) Note that the weight given to the final return Gt(n) is the sum to ∞ of weights from step n onwards, i.e. it is a geometric series. It makes sense to put more weight on the final, actual return. This leads directly to forward-view TD(λ), where we sample trajectories of n steps and update the value function according to: V(St)←V(St)+α(Gtλ−V(St)) Now, the forward view has a shortcoming, which is that we need to wait until we have sampled n steps into the future, before we can update the value function. Thus it suffers similar downside to MC update, where we cannot update the value function immediately after each step. <h2 id="backward-view-td-lambda"><a class="header" href="#backward-view-td-lambda">Backward View TD Lambda</a></h2> One key idea is eligibility traces. In deciding to assign credit to past events for a current reward, there are generally two intuitive heuristics to use: <ul> <li>Frequency heuristic: assign credit to most frequent recent states</li> <li>Recency heuristic: assign credit to most recent states</li> </ul> The eligiblity trace combines both heuristics in a simple formula: <ul> <li>E0(s)=0</li> <li>Et(s)=γλEt−1(s)+1(St=s)</li> </ul> The eligibility trace gives us a weight at a given time step for each state s. This weight tells us how much credit we should assign to s for a reward at the current time step. The Backward View TD Lambda uses this idea: <ul> <li>Keep an eligibility trace for every state s</li> <li>Update value V(s) for every state s in proportion to the TD-error δt and eligibility trace Et(s): δtV(s)=Rt+1+γV(St+1)−V(St)←V(s)+αδtEt(s)</li> </ul> Observe that δt is just our update for TD(0) with a single step look ahead, i.e. Gt(1)−V(St). Thus we can see that when λ=0, only the current state is updated, since Et(s)=1(St=s). This results in the TD(0) update: V(St)←V(St)+αδt. On the other extreme, when λ=1, all credit is deferred until the end of the episode (not sure I see this from the formula). Thus it is equivalent to MC update. In fact, there is a theorem that the sum of offline updates is identical for both forward view and backward view TD-lambda. This is nice because the backward view with eligibility traces makes it easy to implement, as we never need to look forward into the future. We just need to keep track of eligibility traces at each time step, and then apply the update to all states at each step. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="model-free-control"><a class="header" href="#model-free-control">Model Free Control</a></h1> <a href="https://www.youtube.com/watch?v=0g4j2k_Ggc4&list=PLqYmG7hTraZDM-OYHWgPebj2MfCFzFObQ&index=5">Lecture 5: Model Free Control</a> All lectures are building up to this point, to optimize a problem where we do not have access to the underlying MDP. For such problems, we either do not know the underlying MDP, or it is too big to use (e.g. game of Go). On policy vs Off policy <ul> <li>On policy learning is to learn on the job. Learn about policy π based on experience sampled from π</li> <li>Off policy learning is to learn by observing others. Learn about policy π by sampling another robot (or human) experience μ</li> </ul> Start with the simpler case, which is on policy learning. The basic framework is generalized policy iteration (recap), which alternates between: <ul> <li>Policy evaluation: estimate vπ</li> <li>Policy improvement: generate a new π′≥π</li> </ul> Naive case: policy iteration with Monte-Carlo evaluation. Basically, we use MC policy evaluation to update our value function, and then do greedy policy improvement. Would this work? <ul> <li>No. The main problem is that previously when we had access to the underlying MDP, we could do greedy policy improvement because we had access to the transition dynamics. Specifically, when we do policy improvement, we want to compute: π′(s)=aargmaxRsa+Pss′aV(s′)</li> <li>However in model-free control, we do not have access to Pss′a, meaning that we do not know what probabilities determine the state we will end up in given action a. So there is no clear way to do greedy policy improvement if we only have an estimate of V(s)∀s.</li> <li>To deal with this issue, we can do greedy policy improvement over Q(s,a) instead. Then we can simply take: π′(s)=a∈AargmaxQ(s,a)</li> </ul> So now we do generalized policy iteration with action-value function. <ul> <li>Start with Q,π</li> <li>Update action-value function Q=qπ</li> <li>Greedily update policy to π=greedy(Q)</li> </ul> However, we still have another problem, which is the exploration issue. If we act greedily all the time, there is no guarantee that we will explore all states and thus find the optimal policy. <h2 id="toy-example-greedy-action-selection"><a class="header" href="#toy-example-greedy-action-selection">Toy Example: Greedy Action Selection</a></h2> Choose between two doors: <ul> <li>Open left door: reward <code>0</code>. V(left)=0</li> <li>Open right door: reward <code>+1</code>. V(right)=+1</li> <li>Open right door: reward <code>+3</code>. V(right)=+2</li> <li>Open right door: reward <code>+2</code>. V(right)=+2</li> </ul> The greedy policy will lock us onto right door forever. But we will never know if the left door actually has higher mean return. <h2 id="ϵ-greedy-exploration"><a class="header" href="#ϵ-greedy-exploration">ϵ-Greedy Exploration</a></h2> The simplest idea for ensuring continual exploration. <ul> <li>Try all m actions with non-zero probability</li> <li>With probability 1−ϵ choose the greedy action</li> <li>With probability ϵ choose an action at random</li> </ul> π(a∣s)={ϵ/m+1−ϵ ϵ/m if a∗=argmaxa∈AQ(s,a)otherwise Note that ϵ/m is added for the first case as well since the action chosen at random can include the greedy policy a∗ as well. ϵ-greedy policy is important because there is a theorem to assure us that we will indeed get a policy improvement on every step. <blockquote> Theorem. For any ϵ-greedy policy π, the ϵ-greedy policy π′ with respect to qπ is an improvement, i.e. vπ′(s)≥vπ(s) Proof. qπ(s,π′(s))=a∈A∑π′(a∣s)qπ(s,a)=ϵ/ma∈A∑qπ(s,a)+(1−ϵ)a∈Amaxqπ(s,a)≥ϵ/ma∈A∑qπ(s,a)+(1−ϵ)a∈A∑1−ϵπ(a∣s)−ϵ/mqπ(s,a)=a∈A∑π(a∣s)qπ(s,a)=vπ(s) Therefore from the policy improvement theorem, vπ′(s)≥vπ(s). </blockquote> The key step in the proof is the transition from line 2 to line 3. The idea is that the maximum q-value (by choosing the greedy action) will be greater than or equal to any weighted average of qπ(s,a). Hence we choose a clever weighted average such that we can end up with ∑a∈Aπ(a∣s)qπ(s,a) in line 4. Note that it is indeed a weighted average because of the following. Note that π(a∣s) must sum to 1 over all actions as it is a valid policy. And since there are m unique actions, we multiply the constant ϵ/m by m. a∈A∑1−ϵπ(a∣s)−ϵ/m=1−ϵ1a∈A∑[π(a∣s)−ϵ/m]=1−ϵ1[1−m⋅ϵ/m]=1 An idea that we encountered earlier. We do not need to fully evaluate the policy before we do a greedy improvement. In the context of Monte Carlo policy evaluation, in the extreme case, we can update the policy after every episode instead of gathering many episodes. How can we guarantee that we find the optimal policy π∗? We need to ensure that our algorithm balances two things: (i) suitably explore all options and (ii) ensure that at the end, we converge on a greedy policy. This leads us to GLIE, which is a property that we want our algo to have. <blockquote> Definition Greedy in the Limit with Infinite Exploration (GLIE). <ul> <li>All state-action pairs are explored infinitely many times, i.e. k→∞limNk(s,a)=∞</li> <li>The policy converges on a greedy policy, i.e. k→∞limπk(a∣s)=1(a=a′∈AargmaxQk(s,a′))</li> </ul> </blockquote> One simple way to get GLIE is to use ϵ-greedy with a decaying schedule for ϵk=k1. <h2 id="glie-monte-carlo-control"><a class="header" href="#glie-monte-carlo-control">GLIE Monte Carlo Control</a></h2> This brings us to GLIE Monte Carlo control. <blockquote> Algorithm GLIE Monte-Carlo Control. <ul> <li>Sample kth episode using policy π:{S1,A1,R2,...,ST}∼π</li> <li>For each state St and action At in the episode, update N(St,At)Q(St,At)←N(St,At)+1←Q(St,At)+N(St,At)1(Gt−Q(St,At))</li> <li>Improve policy based on the new action-value function: ϵπ←1/k←ϵ-greedy(Q)</li> </ul> </blockquote> <h2 id="mc-vs-td-control"><a class="header" href="#mc-vs-td-control">MC vs TD Control</a></h2> <ul> <li>TD learning has several advantages over MC: <ul> <li>Lower variance</li> <li>Online</li> <li>Can deal with incomplete sequences</li> </ul> </li> <li>Natural idea: use TD instead of MC in our control loop <ul> <li>Apply TD to Q(S, A)</li> <li>Use ϵ-greedy policy improvement</li> <li>Update every time step</li> <li>This is probably the most well known RL algorithm (Sarsa)</li> </ul> </li> </ul> Sarsa policy evaluation update step: Q(S,A)←Q(S,A)+α(R+γQ(S′,A′)−Q(S,A)) Note that we are updating the Q value for one single state-action pair. We take action A on state S and observe reward R, and use that to update the Q-value. In addition, we also sample a next action A′ and corresponding resultant state S′, and we bootstrap the Q-value to use Q(S′,A′) to also update the Q-value. So it corresponds to a one-step lookahead in TD. So the off-policy control with Sarsa algo. For every time step: <ul> <li>Policy evaluation with Sarsa: Q=qπ</li> <li>Policy improvement using ϵ-greedy</li> </ul> <blockquote> Algorithm. Sarsa algorithm for on-policy control. <ul> <li>Initialize Q(s,a),∀s∈S,a∈A(s) arbitrarily</li> <li>Repeat (for each episode): <ul> <li>Initialize S</li> <li>Choose A from S using policy derived from Q (e.g. ϵ-greedy choice)</li> <li>Repeat (for each step of episode): <ul> <li>Take action A, observe R,S′</li> <li>Choose A′ from S′ using policy derived from Q (e.g. ϵ-greedy)</li> <li>Update Q(S,A)←Q(S,A)+α[R+γQ(S′,A′)−Q(S,A)]</li> <li>S←S′, A←A′</li> </ul> </li> <li>Until S is terminal</li> </ul> </li> </ul> </blockquote> Note that this is a fundamentally on-policy algorithm, because the A′,S′ that we sample and use to bootstrap is also the next action and state we end up in. <blockquote> Algorithm. Sarsa converges to the optimal action value function, Q(s,a)→q∗(s,a) under the following conditions: <ul> <li>GLIE sequence of policies πt(a∣s)</li> <li>Robbins Monro sequence of step sizes αt: t=1∑∞αtt=1∑∞αt2=∞<∞</li> </ul> </blockquote> <h2 id="n-step-sarsa"><a class="header" href="#n-step-sarsa">n-step Sarsa</a></h2> As before, we saw that n-step algorithm gets the best of both worlds in betwen MC and TD. So we do the same here. Consider the following n-step returns for n=1,2,∞: <ul> <li>n=1, qt(1)=Rt+1+γQ(St+1)</li> <li>n=2, qt(2)=Rt+1+γRt+2+γ2Q(St+2)</li> <li>n=∞, qt(∞)=Rt+1+γRt+2+...+γT−1Rt</li> </ul> Define the n-step Q-return: qt(n)=Rt+1+γRt+2+...+γn−1Rt+n+γnQ(St+n) Sarsa update Q(s,a) towards the n-step Q-return: Q(St,At)←Q(St,At)+α(qt(n)−Q(St,At)) <h2 id="forward-view-sarsaλ"><a class="header" href="#forward-view-sarsaλ">Forward View Sarsa(λ)</a></h2> As before, we saw that the n-step return itself is noisy and sensitive to hyperparameter choice of n and α. So the better way is to average the value over all n steps. <ul> <li>The qλ return combines all n-step Q-returns qt(n)</li> <li>Using weight (1−λ)λn−1, we have: qtλ=(1−λ)n=1∑∞λn−1qt(n)</li> <li>And the forward view Sarsa(λ) is: Q(St,At)←Q(St,At)+α(qtλ−Q(St,At))</li> </ul> <h2 id="backward-view-sarsaλ"><a class="header" href="#backward-view-sarsaλ">Backward View Sarsa(λ)</a></h2> Recall that we used eligibility traces to construct the backward view TD(λ). As the forward view algo is not an online policy - we need to wait until the end of the episode to do the update. <ul> <li>Just like TD(λ), we use eligibility traces in an online algorithm</li> <li>But Sarsa(λ) has one eligibility trace for each state-action pair instead of just for every state E0(s,a)Et(s,a)=0=γλEt−1(s,a)+1(St=s,At=a)</li> <li>Q(s,a) is updated for every state s and action a</li> <li>In proportion to TD-error δt and eligibility trace Et(s,a): δtQ(s,a)=Rt+1+γQ(St+1,At+1)−Q(St,At)←Q(s,a)+αδtEt(s,a)</li> </ul> <blockquote> Algorithm. Sarsa(λ) On Policy Algorithm. <ul> <li>Initialize Q(s,a) arbitrarily, for all s∈S,a∈A</li> <li>Repeat (for each episode): <ul> <li>E(s,a)=0, for all s∈S,a∈A</li> <li>Initialize S,A</li> <li>Repeat (for each step of episode): <ul> <li>Take action A, observe R,S′</li> <li>Choose A′ from S′, using policy derived from Q (E.g. ϵ-greedy)</li> <li>δ←R+γQ(S′,A′)−Q(S,A)</li> <li>E(S,A)←E(S,A)+1</li> <li>For all s∈S,a∈A: <ul> <li>Q(s,a)←Q(s,a)+αδE(s,a)</li> <li>E(s,a)←γλE(s,a)</li> </ul> </li> <li>S←S′,A←A′</li> </ul> </li> <li>Until S is terminal</li> </ul> </li> </ul> </blockquote> Note that for a given step we have a single value of δ which is our TD error, but we propagate that to all s,a pairs based on the eligibility trace, as potentially every s,a pair could have contributed to it. <h2 id="off-policy-learning"><a class="header" href="#off-policy-learning">Off Policy Learning</a></h2> So far we have been looking at on-policy learning. However it is often useful to do off policy learning, i.e. evaluate a target policy π(a∣s) to compute vπ(s) or qπ(s,a), while we follow the behaviour policy μ(a∣s):{S1,A1,R2,...,ST∼μ}. Of course in this case, μ=π. Why is off policy learning useful? <ul> <li>We can learn from observing humans or other agents</li> <li>We can re-use experience that was previously generated from old policies π1,π2,...πt−1, possibly in a batched manner</li> <li>We can learn about the optimal policy while following the exploratory policy</li> <li>We can learn about multiple policies while following one policy</li> </ul> First mechanism is importance sampling. The main idea is to estimate the expectation of a different distribution by re-weighting the distributions: EX∼P[f(X)]=∑P(X)f(X)=∑Q(X)Q(X)P(X)f(X)=EX∼Q[Q(X)P(X)f(X)] We can apply importance sampling to Monte Carlo for Off policy monte carlo learning: <ul> <li>We use returns generated from behaviour policy μ to evaluate π</li> <li>Then we weight the return Gt according to the ratio of probabilities between the two policies</li> <li>We need to apply the correction at every time step along the whole episode, because the change in policy affects every time step Gtπ/μ=μ(At∣St)π(At∣St)μ(At+1∣St+1)π(At+1∣St+1)...μ(AT∣ST)π(AT∣ST)Gt</li> <li>And then update the value towards the corrected return V(St)←V(St)+α(Gtπ/μ−V(St))</li> </ul> While off policy MC learning is theoretically sound, there are some major problems which make it practically useless in practice: <ul> <li>Importance sampling dramatically increases variance, as we are adjusting over every time step, and the cumulative effect over the whole episode makes our estimate of Gtπ/μ vary wildly</li> <li>We also cannot use this adjustment if μ is zero when π is non-zero</li> </ul> So we have to use bootstrapping for importance sampling. This allows us to only adjust the probability for one time step. So we have importance sampling for off policy TD: <ul> <li>We use TD targets generated from μ to evaluate π</li> <li>For TD(0), We weight the TD target R+γV(S′) by importance sampling</li> <li>This means we only need a single importance sampling correction: V(St)←V(St)+α(μ(At∣St)π(At∣St)(Rt+1+γV(St+1))−V(St))</li> <li>This has much lower variance that MC importance sampling, and could work if μ and π do not differ by too much over a single step</li> </ul> As we have seen, importance sampling leads to large variances. The best solution is known as Q-learning, which is specific to TD(0) or Sarsa(0). <ul> <li>Does not require any importance sampling</li> <li>Allows off policy learning of action values Q(s,a)</li> </ul> Recall that μ is the behaviour policy that our agent is actually following, and π is a target policy that we want to learn from. The main idea is that in our Sarsa(0) update step, we update the Q-value towards the target policy π, but allow our agent to continue following the behaviour policy μ. This allows the agent to explore the environment using μ, but learn from the action-value function of π. Specifically: <ul> <li>We choose each next action for the agent using behaviour policy At+1∼μ(.∣St)</li> <li>But use alternative successor action A′∼π(.∣St) in our Q-value update</li> <li>So we update Q(St,At) using A′: Q(St,At)←Q(St,At)+α(Rt+1+γQ(St+1,A′)−Q(St,At))</li> </ul> Note importantly that we are using A′∼π in the Q value above, instead of At+1. This allows us to learn off policy. <h3 id="q-learning-or-sarsa-max"><a class="header" href="#q-learning-or-sarsa-max">Q-Learning (or SARSA-MAX)</a></h3> A special case of Q-learning is the case where the target policy is greedy wrt Q(s,a). This is usually what people refer to as Q-learning. We allow both behaviour and target policies to improve: <ul> <li>The target policy π is greedy wrt Q(s,a), i.e. π(St+1)=a′argmaxQ(St+1,a′)</li> <li>The behaviour policy μ is ϵ-greedy wrt Q(s,a) again</li> <li>The learning target inside the Q-update then simplifies as follows: Rt+1+γQ(St+1,A′)=Rt+1+γQ(St+1,a′argmaxQ(St+1,a′))=Rt+1+a′maxγQ(St+1,a′)</li> </ul> Note that since we are following a greedy target policy, the action chosen will be the Q-maximizing one (line 2). Since we are choosing the Q-maximizing action, we get the maximum Q-value over all possible actions (line 3). This simplifies the equation quite abit, and now it resembles the Bellman optimality equation. This leads us to the well known Q-learning algorithm, which David calls Sarsa-max. The Q-update is: Q(S,A)←Q(S,A)+α(R+γa′maxQ(S′,a′)−Q(S,A)) There is a theorem that tells us that the Q-learning control algorithm converges to the optimal action-value function, i.e. Q(s,a)→q∗(s,a) To wrap up, here is a classification of some algorithms we have so far: <div class="table-wrapper"><table><thead><tr><th style="text-align: center"></th><th style="text-align: center">Full Backup (Dynamic Programming)</th><th style="text-align: center">Sample Backup (Temporal Difference)</th></tr></thead><tbody> <tr><td style="text-align: center">Bellman Expectation Equation for vπ(s)</td><td style="text-align: center">Iterative Policy Evaluation</td><td style="text-align: center">TD Learning</td></tr> <tr><td style="text-align: center">Bellman Expectation Equation for qπ(s,a)</td><td style="text-align: center">Q-policy iteration</td><td style="text-align: center">Sarsa</td></tr> <tr><td style="text-align: center">Bellman Optimality Equation for q∗(s,a)</td><td style="text-align: center">Q-value iteration</td><td style="text-align: center">Q-learning</td></tr> </tbody></table> </div><div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="value-function-approximation"><a class="header" href="#value-function-approximation">Value Function Approximation</a></h1> <a href="https://www.youtube.com/watch?v=UoPei5o4fps&list=PLqYmG7hTraZDM-OYHWgPebj2MfCFzFObQ&index=6">Lecture 6: Value Function Approximation</a> This lecture looks into approximating functions with neural networks to overcome the large state-action space problem. RL often encounters large problems: <ul> <li>Backgammon: 1020 states</li> <li>Go: 10170 states</li> <li>Helicopter: continuous state space</li> </ul> We want to do policy evaluation and control efficiently in large state spaces. So far, we have represented V or Q with a lookup table: <ul> <li>Every state s has an entry V(s)</li> <li>Every state-action pair s,a has an entry Q(s,a)</li> </ul> This is a problem for large MDPs: <ul> <li>Too many states or actions to store in memory</li> <li>It is too slow or data inefficient to learn the value of each state individually</li> </ul> Solution for large MDPs: <ul> <li>Estimate the value function with function approximation using parameters w: <ul> <li>v^(s,w)≈vπ(s)</li> <li>q^(s,a,w)≈qπ(s,a)</li> </ul> </li> <li>Generalizes from seen states to unseen states</li> <li>Update parameters of our function using MC or TD learning</li> </ul> Types of value function approximation (different architectures): <ul> <li>Represent a given state s with some parameters w. Then neural network spits out v^(s,w), which is our value function for being in state s</li> <li>Have a neural network q^(s,a,w), which takes in a state-action pair and spits out the Q value</li> <li>Sometimes, it is more efficient to have a neural network q^(s,w), such that we feed in a single state and we get Q-values for every possible action in a single forward pass, i.e. we get q^(s,a1,w),q^(s,a2,w),...</li> </ul> Which function approximator? We focus on differentiable function approximators that we can easily optimize, i.e. Linear combinations of features, neural networks. Furthermore, we want a training algorithm for a non-iid, non-stationary set of data, so it is not standard supervised learning. <h2 id="incremental-methods"><a class="header" href="#incremental-methods">Incremental Methods</a></h2> <h3 id="gradient-descent"><a class="header" href="#gradient-descent">Gradient Descent</a></h3> Starting with gradient descent. <ul> <li>Let J(w) be a differentiable function of parameter vector w</li> <li>Define the gradient of J(w) to be a vector ∇wJ(w) where ∇wJ(w)[0] is ∂w1∂J(w)</li> <li>To find the local minimum of J(w), we adjust the parameter w in the -ve gradient direction: △w=−21α∇wJ(w)</li> </ul> Goal: find parameter vector w minimizing mean squared error between approximate value fn v^(s,w) and true oracle value fn vπ(s) (assuming we know the oracle) J(w)=Eπ[(vπ(S)−v^(S,w)2)] Gradient descent finds a local minimum: △w=−21α∇wJ(w)=αEπ[(vπ(S)−v^(S,w))∇wv^(S,w)] Stochastic gradient descent samples the gradient: △w=α(vπ(S)−v^(S,w))∇wv^(S,w) The nice thing about SGD is that it still converges under non-stationary environment. The expected update is equal to full gradient update. <h3 id="feature-vectors"><a class="header" href="#feature-vectors">Feature Vectors</a></h3> To represent a state, we use a feature vector. x(S)=<svg xmlns="http://www.w3.org/2000/svg" width='0.875em' height='4.200em' viewBox='0 0 875 4200'><path d='M863,9c0,-2,-2,-5,-6,-9c0,0,-17,0,-17,0c-12.7,0,-19.3,0.3,-20,1 c-5.3,5.3,-10.3,11,-15,17c-242.7,294.7,-395.3,682,-458,1162c-21.3,163.3,-33.3,349, -36,557 l0,684c0.2,6,0,26,0,60c2,159.3,10,310.7,24,454c53.3,528,210, 949.7,470,1265c4.7,6,9.7,11.7,15,17c0.7,0.7,7,1,19,1c0,0,18,0,18,0c4,-4,6,-7,6,-9 c0,-2.7,-3.3,-8.7,-10,-18c-135.3,-192.7,-235.5,-414.3,-300.5,-665c-65,-250.7,-102.5, -544.7,-112.5,-882c-2,-104,-3,-167,-3,-189 l0,-692c0,-162.7,5.7,-314,17,-454c20.7,-272,63.7,-513,129,-723c65.3, -210,155.3,-396.3,270,-559c6.7,-9.3,10,-15.3,10,-18z'/></svg>x1(S)⋮xn(S)<svg xmlns="http://www.w3.org/2000/svg" width='0.875em' height='4.200em' viewBox='0 0 875 4200'><path d='M76,0c-16.7,0,-25,3,-25,9c0,2,2,6.3,6,13c21.3,28.7,42.3,60.3, 63,95c96.7,156.7,172.8,332.5,228.5,527.5c55.7,195,92.8,416.5,111.5,664.5 c11.3,139.3,17,290.7,17,454c0,28,1.7,43,3.3,45l0,609 c-3,4,-3.3,16.7,-3.3,38c0,162,-5.7,313.7,-17,455c-18.7,248,-55.8,469.3,-111.5,664 c-55.7,194.7,-131.8,370.3,-228.5,527c-20.7,34.7,-41.7,66.3,-63,95c-2,3.3,-4,7,-6,11 c0,7.3,5.7,11,17,11c0,0,11,0,11,0c9.3,0,14.3,-0.3,15,-1c5.3,-5.3,10.3,-11,15,-17 c242.7,-294.7,395.3,-681.7,458,-1161c21.3,-164.7,33.3,-350.7,36,-558 l0,-744c-2,-159.3,-10,-310.7,-24,-454c-53.3,-528,-210,-949.7, -470,-1265c-4.7,-6,-9.7,-11.7,-15,-17c-0.7,-0.7,-6.7,-1,-18,-1z'/></svg> For example, the features (numeric) could be: <ul> <li>Distance of robot to landmarks</li> <li>Trends in the stock market</li> <li>Configuration of pawn on a chess board</li> </ul> <h3 id="linear-value-function-approximation"><a class="header" href="#linear-value-function-approximation">Linear Value Function Approximation</a></h3> Let us represent the value function using a linear combination of features (i.e. just a dot product between two vectors): v^(S,w)=x(S)Tw=j=1∑nxj(S)wj The nice thing is that linear approximator is quadratic in the parameters w, so it is a convex optimization problem, i.e. SGD will converge on the global optimum: J(w)=Eπ[(vπ(S)−x(S)Tw)2] The gradient update is really simple: ∇wv^(S,w)△w=x(S)=α(vπ(S)−v^(S,w))x(S) Note that we are just subbing the simple expression for ∇wv^(S,w) into the general △w formula above. The update may be interpreted as <code>step-size x prediction error x feature value</code>. This means that features with high correlation with the prediction error will have large gradient updates intuitively. We can think of table lookup as a special case of linear value function approximation. Suppose we use a table lookup feature (1-hot) as follows: xtable(S)=<svg xmlns="http://www.w3.org/2000/svg" width='0.875em' height='4.200em' viewBox='0 0 875 4200'><path d='M863,9c0,-2,-2,-5,-6,-9c0,0,-17,0,-17,0c-12.7,0,-19.3,0.3,-20,1 c-5.3,5.3,-10.3,11,-15,17c-242.7,294.7,-395.3,682,-458,1162c-21.3,163.3,-33.3,349, -36,557 l0,684c0.2,6,0,26,0,60c2,159.3,10,310.7,24,454c53.3,528,210, 949.7,470,1265c4.7,6,9.7,11.7,15,17c0.7,0.7,7,1,19,1c0,0,18,0,18,0c4,-4,6,-7,6,-9 c0,-2.7,-3.3,-8.7,-10,-18c-135.3,-192.7,-235.5,-414.3,-300.5,-665c-65,-250.7,-102.5, -544.7,-112.5,-882c-2,-104,-3,-167,-3,-189 l0,-692c0,-162.7,5.7,-314,17,-454c20.7,-272,63.7,-513,129,-723c65.3, -210,155.3,-396.3,270,-559c6.7,-9.3,10,-15.3,10,-18z'/></svg>1(S=s1)⋮1(S=sn)<svg xmlns="http://www.w3.org/2000/svg" width='0.875em' height='4.200em' viewBox='0 0 875 4200'><path d='M76,0c-16.7,0,-25,3,-25,9c0,2,2,6.3,6,13c21.3,28.7,42.3,60.3, 63,95c96.7,156.7,172.8,332.5,228.5,527.5c55.7,195,92.8,416.5,111.5,664.5 c11.3,139.3,17,290.7,17,454c0,28,1.7,43,3.3,45l0,609 c-3,4,-3.3,16.7,-3.3,38c0,162,-5.7,313.7,-17,455c-18.7,248,-55.8,469.3,-111.5,664 c-55.7,194.7,-131.8,370.3,-228.5,527c-20.7,34.7,-41.7,66.3,-63,95c-2,3.3,-4,7,-6,11 c0,7.3,5.7,11,17,11c0,0,11,0,11,0c9.3,0,14.3,-0.3,15,-1c5.3,-5.3,10.3,-11,15,-17 c242.7,-294.7,395.3,-681.7,458,-1161c21.3,-164.7,33.3,-350.7,36,-558 l0,-744c-2,-159.3,-10,-310.7,-24,-454c-53.3,-528,-210,-949.7, -470,-1265c-4.7,-6,-9.7,-11.7,-15,-17c-0.7,-0.7,-6.7,-1,-18,-1z'/></svg> And suppose we have a parameter vector of size n, such that we have one parameter for each state. Then we have: v^(S,w)=<svg xmlns="http://www.w3.org/2000/svg" width='0.875em' height='4.200em' viewBox='0 0 875 4200'><path d='M863,9c0,-2,-2,-5,-6,-9c0,0,-17,0,-17,0c-12.7,0,-19.3,0.3,-20,1 c-5.3,5.3,-10.3,11,-15,17c-242.7,294.7,-395.3,682,-458,1162c-21.3,163.3,-33.3,349, -36,557 l0,684c0.2,6,0,26,0,60c2,159.3,10,310.7,24,454c53.3,528,210, 949.7,470,1265c4.7,6,9.7,11.7,15,17c0.7,0.7,7,1,19,1c0,0,18,0,18,0c4,-4,6,-7,6,-9 c0,-2.7,-3.3,-8.7,-10,-18c-135.3,-192.7,-235.5,-414.3,-300.5,-665c-65,-250.7,-102.5, -544.7,-112.5,-882c-2,-104,-3,-167,-3,-189 l0,-692c0,-162.7,5.7,-314,17,-454c20.7,-272,63.7,-513,129,-723c65.3, -210,155.3,-396.3,270,-559c6.7,-9.3,10,-15.3,10,-18z'/></svg>1(S=s1)⋮1(S=sn)<svg xmlns="http://www.w3.org/2000/svg" width='0.875em' height='4.200em' viewBox='0 0 875 4200'><path d='M76,0c-16.7,0,-25,3,-25,9c0,2,2,6.3,6,13c21.3,28.7,42.3,60.3, 63,95c96.7,156.7,172.8,332.5,228.5,527.5c55.7,195,92.8,416.5,111.5,664.5 c11.3,139.3,17,290.7,17,454c0,28,1.7,43,3.3,45l0,609 c-3,4,-3.3,16.7,-3.3,38c0,162,-5.7,313.7,-17,455c-18.7,248,-55.8,469.3,-111.5,664 c-55.7,194.7,-131.8,370.3,-228.5,527c-20.7,34.7,-41.7,66.3,-63,95c-2,3.3,-4,7,-6,11 c0,7.3,5.7,11,17,11c0,0,11,0,11,0c9.3,0,14.3,-0.3,15,-1c5.3,-5.3,10.3,-11,15,-17 c242.7,-294.7,395.3,-681.7,458,-1161c21.3,-164.7,33.3,-350.7,36,-558 l0,-744c-2,-159.3,-10,-310.7,-24,-454c-53.3,-528,-210,-949.7, -470,-1265c-4.7,-6,-9.7,-11.7,-15,-17c-0.7,-0.7,-6.7,-1,-18,-1z'/></svg>⋅<svg xmlns="http://www.w3.org/2000/svg" width='0.875em' height='4.200em' viewBox='0 0 875 4200'><path d='M863,9c0,-2,-2,-5,-6,-9c0,0,-17,0,-17,0c-12.7,0,-19.3,0.3,-20,1 c-5.3,5.3,-10.3,11,-15,17c-242.7,294.7,-395.3,682,-458,1162c-21.3,163.3,-33.3,349, -36,557 l0,684c0.2,6,0,26,0,60c2,159.3,10,310.7,24,454c53.3,528,210, 949.7,470,1265c4.7,6,9.7,11.7,15,17c0.7,0.7,7,1,19,1c0,0,18,0,18,0c4,-4,6,-7,6,-9 c0,-2.7,-3.3,-8.7,-10,-18c-135.3,-192.7,-235.5,-414.3,-300.5,-665c-65,-250.7,-102.5, -544.7,-112.5,-882c-2,-104,-3,-167,-3,-189 l0,-692c0,-162.7,5.7,-314,17,-454c20.7,-272,63.7,-513,129,-723c65.3, -210,155.3,-396.3,270,-559c6.7,-9.3,10,-15.3,10,-18z'/></svg>w1⋮wn<svg xmlns="http://www.w3.org/2000/svg" width='0.875em' height='4.200em' viewBox='0 0 875 4200'><path d='M76,0c-16.7,0,-25,3,-25,9c0,2,2,6.3,6,13c21.3,28.7,42.3,60.3, 63,95c96.7,156.7,172.8,332.5,228.5,527.5c55.7,195,92.8,416.5,111.5,664.5 c11.3,139.3,17,290.7,17,454c0,28,1.7,43,3.3,45l0,609 c-3,4,-3.3,16.7,-3.3,38c0,162,-5.7,313.7,-17,455c-18.7,248,-55.8,469.3,-111.5,664 c-55.7,194.7,-131.8,370.3,-228.5,527c-20.7,34.7,-41.7,66.3,-63,95c-2,3.3,-4,7,-6,11 c0,7.3,5.7,11,17,11c0,0,11,0,11,0c9.3,0,14.3,-0.3,15,-1c5.3,-5.3,10.3,-11,15,-17 c242.7,-294.7,395.3,-681.7,458,-1161c21.3,-164.7,33.3,-350.7,36,-558 l0,-744c-2,-159.3,-10,-310.7,-24,-454c-53.3,-528,-210,-949.7, -470,-1265c-4.7,-6,-9.7,-11.7,-15,-17c-0.7,-0.7,-6.7,-1,-18,-1z'/></svg> And we can see that this reduces to a table lookup where the parameter wj represents the state value for each state j. <h3 id="estimating-the-oracle"><a class="header" href="#estimating-the-oracle">Estimating the Oracle</a></h3> So far, we have assumed that the true oracle value function vπ(s) is available, but in RL there is no true label, only rewards. So in practice, we need to substitute a target for vπ(s): <ul> <li>For MC, the target is the return Gt: △w=α(Gt−v^(St,w))∇wv^(St,w)</li> <li>For TD(0), the target is the TD target Rt+1+γv^(St+1,w): △w=α(Rt+1+γv^(St+1,w)−v^(St,w))∇wv^(St,w)</li> <li>For TD(λ), the target is the λ-return Gtλ: △w=α(Gtλ−v^(St,w))∇wv^(St,w)</li> </ul> <h3 id="monte-carlo-with-value-function-approximation"><a class="header" href="#monte-carlo-with-value-function-approximation">Monte Carlo with Value Function Approximation</a></h3> We can think of our algorithm as supervised learning. <ul> <li>Treat the return Gt as an unbiased noisy sample of the true value vπ(St)</li> <li>We therefore are applying supervised learning to "training data": <St,G1>,<S2,G2>,...,<ST,GT></li> <li>For example, using linear MC policy evaluation: △w=α(Gt−v^(St,w))∇wv^(St,w)=α(Gt−v^(St,w))x(St)</li> <li>MC evaluation converges to a local optimum even when using non-linear value function approximation</li> </ul> <h3 id="td-with-value-function-approximation"><a class="header" href="#td-with-value-function-approximation">TD with Value Function Approximation</a></h3> The same applies to TD learning, but we have some biased estimate: <ul> <li>The TD target Rt+1+γv^(St+1,w) is a biased sample of the true value vπ(St) - it's biased because our own value function is a biased estimate</li> <li>We can still apply supervised learning to the "training data": <S1,R2+γv^(S2,w)>,<S2,R3+γv^(S3,w)>,...,<St−1,RT></li> <li>For example using linear TD(0): △w=α(R+γv^(S′,w)−v^(S,w))∇wv^(S,w)=αδx(S)</li> <li>There is a theorem showing that for linear TD(0), despite the bias, we will always converge (close) to the global optimum</li> </ul> <blockquote> Note: There is a little inconsistency in the above formula, once we start introducing bootstrapped approximations of the return. Recall that when we used the oracle vπ to represent the target and took the derivative, only ∇wv^(S,w) enters the derivative as we treat the oracle value as a constant. However, once we introduce v^ itself to substitute the oracle function, we should technically include that term in the derivative as well. As it turns out, this is not a good idea and will not lead to convergence. There is some theoretical analysis for this to justify it. </blockquote> <h3 id="tdλ-with-value-function-approximation"><a class="header" href="#tdλ-with-value-function-approximation">TD(λ) with Value Function Approximation</a></h3> And again, we can do the same with TD-λ, since the λ-return Gtλ is also a biased sample of the true value vπ(s): <ul> <li>The training data is now: <S1,G1λ>,<S2,G2λ>,...,<ST−1,GT−1λ></li> <li>The forward view linear TD(λ) is: △w=α(Gtλ−v^(St,w))∇wv^(St,w)=α(Gtλ−v^(St,w))x(St)</li> <li>The backward view linear TD(λ) is: δtEt△w=Rt+1+γv^(St+1,w)−v^(St,w)=γλEt−1+x(St)=αδtEt</li> <li>There is a theorem to show that the forward view and backward view linear TD(λ) are equivalent.</li> </ul> For the backward view, notice that the eligibility trace is now updated using the gradient wrt the parameter vector, namely ∇wv^(St,w), which is of the same dimensionality as w. More precisely, the eligibility trace is the decaying accumulation of past gradients. In the linear case, this is an accumulation of the feature vector x(St). It is a bit unintuitive to understand why we use the accumulated gradient as the eligibility trace, but I suppose it is proved in the equivalence theorem between the forward and backward view. Perhaps we can just think of it as "the features which we see the most often will have high eligibility trace". <h2 id="control-with-value-function-approximation"><a class="header" href="#control-with-value-function-approximation">Control with Value Function Approximation</a></h2> <ul> <li>Start with some random parameter vector w</li> <li>Set policy based on some greedy function π=ϵ-greedy(qw)</li> <li>Do policy evaluation q^(.,.,w)≈qπ</li> </ul> First we need to do everything again wrt to action-value function instead of value function to perform this algorithm. The steps are: <ul> <li>Approximate the action-value function q^(S,A,w)≈qπ(S,A)</li> <li>Minimize the mean squared error between approximate action value function and true oracle action value qπ(S,A): J(w)=Eπ[(qπ(S,A)−q^(S,A,w)2)]</li> <li>Use SGD to find a local minimum: −21∇wJ(w)△w=(qπ(S,A)−q^(S,A,w))∇wq^(S,A,w)=α(qπ(S,A)−q^(S,A,w))∇wq^(S,A,w)</li> <li>Again, we represent the state and action by a feature vector: x(S,A)=<svg xmlns="http://www.w3.org/2000/svg" width='0.875em' height='4.200em' viewBox='0 0 875 4200'><path d='M863,9c0,-2,-2,-5,-6,-9c0,0,-17,0,-17,0c-12.7,0,-19.3,0.3,-20,1 c-5.3,5.3,-10.3,11,-15,17c-242.7,294.7,-395.3,682,-458,1162c-21.3,163.3,-33.3,349, -36,557 l0,684c0.2,6,0,26,0,60c2,159.3,10,310.7,24,454c53.3,528,210, 949.7,470,1265c4.7,6,9.7,11.7,15,17c0.7,0.7,7,1,19,1c0,0,18,0,18,0c4,-4,6,-7,6,-9 c0,-2.7,-3.3,-8.7,-10,-18c-135.3,-192.7,-235.5,-414.3,-300.5,-665c-65,-250.7,-102.5, -544.7,-112.5,-882c-2,-104,-3,-167,-3,-189 l0,-692c0,-162.7,5.7,-314,17,-454c20.7,-272,63.7,-513,129,-723c65.3, -210,155.3,-396.3,270,-559c6.7,-9.3,10,-15.3,10,-18z'/></svg>x1(S,A)⋮xn(S,A)<svg xmlns="http://www.w3.org/2000/svg" width='0.875em' height='4.200em' viewBox='0 0 875 4200'><path d='M76,0c-16.7,0,-25,3,-25,9c0,2,2,6.3,6,13c21.3,28.7,42.3,60.3, 63,95c96.7,156.7,172.8,332.5,228.5,527.5c55.7,195,92.8,416.5,111.5,664.5 c11.3,139.3,17,290.7,17,454c0,28,1.7,43,3.3,45l0,609 c-3,4,-3.3,16.7,-3.3,38c0,162,-5.7,313.7,-17,455c-18.7,248,-55.8,469.3,-111.5,664 c-55.7,194.7,-131.8,370.3,-228.5,527c-20.7,34.7,-41.7,66.3,-63,95c-2,3.3,-4,7,-6,11 c0,7.3,5.7,11,17,11c0,0,11,0,11,0c9.3,0,14.3,-0.3,15,-1c5.3,-5.3,10.3,-11,15,-17 c242.7,-294.7,395.3,-681.7,458,-1161c21.3,-164.7,33.3,-350.7,36,-558 l0,-744c-2,-159.3,-10,-310.7,-24,-454c-53.3,-528,-210,-949.7, -470,-1265c-4.7,-6,-9.7,-11.7,-15,-17c-0.7,-0.7,-6.7,-1,-18,-1z'/></svg></li> <li>Represent action value function by a linear combination of features: q^(S,A,w)=x(S,A)Tw=j=1∑nxj(S,A)wj</li> <li>Do an SGD update: ∇wq^(S,A,w)△w=x(S,A)=α(qπ(S,A)−q^(S,A,w))x(S,A)</li> </ul> <h3 id="incremental-control-algorithms"><a class="header" href="#incremental-control-algorithms">Incremental Control Algorithms</a></h3> Like prediction, we need to substitute a target for the unknown oracle qπ(S,A). We sub out all the vπ for an approximate target: <ul> <li>For MC, target is the return Gt △w=α(Gt−q^(St,At,w))∇wq^(St,At,w)</li> <li>For TD(0), the target is the TD target Rt+1+γQ(St+1,At+1):: △w=α(Rt+1+γq^(St+1,At+1,w))−q^(St,At,w)∇wq^(St,At,w)</li> <li>For forward view TD(λ), the target is the action value λ-return: △w=α(qtλ−q^(St,At,w))∇wq^(St,At,w)</li> <li>For backward view TD(λ), the equivalent update is: δtEt△w=Rt+1+γq^(St+1,At+1,w)−q^(St,At,w)=γλEt−1+∇wq^(St,At,w)=αδtEt</li> </ul> Should we bootstrap? Empirically across many examples, we almost always have the case that: <ul> <li>MC takes too many steps because variance is too high</li> <li>TD(0) always has a large efficiency gain compared to MC</li> <li>There's always some λ value in between which is better than TD(0)</li> </ul> <h2 id="batch-methods"><a class="header" href="#batch-methods">Batch Methods</a></h2> Motivation: <ul> <li>Gradient descent is simple and appealing</li> <li>But it is not sample efficient (we throw a sample away as soon as we use it once)</li> <li>Batch methods seek to find the best fitting value function, given the agent's experience ("training data")</li> </ul> <h3 id="least-squares-prediction"><a class="header" href="#least-squares-prediction">Least Squares Prediction</a></h3> The problem becomes the following: <ul> <li>Given our value function approximation v^(s,w)≈vπ(s)</li> <li>And experience D consisting of <code><state, value></code> pairs</li> </ul> D={<s1,v1π>,<s2,v2π>,...,<sT,vTπ>,} <ul> <li>Find the parameters w that give the best fitting function v^(s,w)</li> </ul> Least squares algorithms simply try to find w that minimizes the sum of squares error between v^(st,w) and target values vtπ: LS(w)=t=1∑T[vtπ−v^(st,w)]2=ED[vπ−v^(s,w)]2 <h3 id="sgd-with-experience-replay"><a class="header" href="#sgd-with-experience-replay">SGD with Experience Replay</a></h3> It turns out there is a really easy way to find the least squares solution, using experience replay. The idea is to just keep using the data over and over again, instead of throwing away every sample after each update. Given experience comprising of: D={⟨s1,v1π,⟩⟨s2,v2π,⟩...,⟨sT,vTπ,⟩} Repeat: <ol> <li>Sample state, value from experience: ⟨s,vπ⟩∼D</li> <li>Apply SGD update: △w=α(vπ−v^(s,w))∇wv^(s,w)</li> </ol> It can be shown that this converges to the least squares solution: wπ=wargminLS(w) <h3 id="experience-replay-in-deep-q-networks-dqn"><a class="header" href="#experience-replay-in-deep-q-networks-dqn">Experience Replay in Deep Q-Networks (DQN)</a></h3> DQN (for atari games) uses experience replay and fixed Q-targets: <ul> <li>Take action at according to ϵ-greedy policy</li> <li>Store transition (st,at,rt+1,st+1) in replay memory D</li> <li>Sample random mini batch of transitions (s,a,r,s′) <ul> <li>Small batch size of <code>64</code> is sufficient</li> </ul> </li> <li>Maintain two neural networks that estimate Q-values: <ul> <li>The old reference neural network is frozen periodically and used as the target</li> <li>Call its parameters w−</li> <li>The actual neural network we are training has parameters w</li> </ul> </li> <li>Compute Q-learning targets wrt old, fixed parameters w−</li> <li>Optimize MSE between reference Q-network and Q-learning targets: Li(wi)=Es,a,r,s′∼Di[(r+γa′maxQ(s′,a′;wi−)−Q(s,a;wi))2]</li> <li>This is essentially Q-learning with a one-step look ahead, but using the reference network instead of the current active network under training</li> <li>Success of this method depends on its stability in training: <ol> <li>Experience replay helps to stabilize training as it randomly samples from past experience instead of getting batches of highly correlated data</li> <li>Fixed Q-targets - fixing the reference neural network helps to stabilize the targets and thus training</li> </ol> </li> <li>The neural network is just a large convolutional neural network <ul> <li>Input state s is a stack of raw pixels from last 4 frames</li> <li>Output is Q(s,a) for <code>18</code> joystick / button positions</li> <li>Reward is the change in score for that step</li> <li>Applied to a large number of Atari games</li> </ul> </li> </ul> <h3 id="linear-least-squares-prediction"><a class="header" href="#linear-least-squares-prediction">Linear Least Squares Prediction</a></h3> Experience replay finds the least squares solution, but it takes many iterations. If we use a linear value function approximation, we can solve the least squares solution directly. At the minimum of LS(w), the expected update must be zero: ED[△w]=0 So the expected update is zero: t=1∑Tx(st)(vtπ−x(st)Tw)=0 Solving for w: t=1∑Tx(st)vtπw=t=1∑Tx(st)x(st)Tw=(t=1∑Tx(st)x(st)T)−1t=1∑Tx(st)vtπ <ul> <li>Note that the matrix inverse is performed on a matrix of size ∣w∣2, where ∣w∣ is the size of the feature / parameter vector. Hence if the number of parameters is small, this is acceptable to take the O(N3) complexity</li> <li>Using Shermann-Morrison, the solution time is reduced to O(N2)</li> </ul> Linear least squares prediction algorithms actually have better convergence properties. <h3 id="least-squares-policy-iteration"><a class="header" href="#least-squares-policy-iteration">Least Squares Policy Iteration</a></h3> <ul> <li>Policy evaluation is done using least squares Q-learning (linear or otherwise)</li> <li>Policy improvement is done using greedy policy improvement as per normal</li> </ul> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="policy-gradient"><a class="header" href="#policy-gradient">Policy Gradient</a></h1> <a href="https://www.youtube.com/watch?v=KHZVXao4qXs&list=PLqYmG7hTraZDM-OYHWgPebj2MfCFzFObQ&index=7">Lecture 7: Policy Gradient</a> Look at methods that update the policy directly, instead of working with value / action-value functions. <ul> <li>In the last lecture we approximated the value of action value function using parameters θ. Vθ(s)Qθ(s,a)≈Vπ(s)≈Qπ(s,a)</li> <li>We generated a policy directly from the value function, e.g. using ϵ-greedy algorithm</li> <li>In this lecture, we will directly parametrise the policy: πθ(s,a)=P[a∣s,θ]</li> <li>Again, we will focus on model free reinforcement learning</li> </ul> <h2 id="value-based-vs-policy-based-rl"><a class="header" href="#value-based-vs-policy-based-rl">Value-based vs Policy-based RL</a></h2> <ul> <li>Value Based <ul> <li>Learn the value function</li> <li>Policy is implicit from the value function (e.g. ϵ-greedy)</li> </ul> </li> <li>Policy Based <ul> <li>No value function</li> <li>Learn policy directly</li> </ul> </li> <li>Actor-Critic <ul> <li>Learn a value function</li> <li>Also learn a policy</li> </ul> </li> </ul> Advantages of Policy-based RL <ul> <li>Better convergence properties than value based <ul> <li>Value based sometimes swings or chatters around the optimum and do not converge</li> </ul> </li> <li>Effective in high dimensional or continuous action spaces <ul> <li>We do not need to compute max action over Q-values</li> <li>E.g. if the action space is continuous, the maximization is not at all trivial and may be prohibitive</li> <li>This may be the main impediment to value-based RL</li> </ul> </li> <li>Can learn stochastic policies</li> </ul> Disadvantages of Policy-base RL <ul> <li>Typically converges to a local rather than global optimum</li> <li>Evaluating a policy is typically inefficient and high variance</li> </ul> Why might we want to have a stochastic policy? <ul> <li>e.g. Rock paper scissors</li> <li>Having a deterministic policy is easily exploited</li> <li>A uniform random policy is optimal</li> </ul> Stochastic policy is also necessary in a case of state-aliasing, in which our state representation cannot exhaustively differentiate states from each other. In this case, we do not have a Markov Decision Process, and there may not exist a deterministic policy that is optimal. That is why we need a stochastic policy. <h2 id="policy-objective-functions"><a class="header" href="#policy-objective-functions">Policy Objective Functions</a></h2> Our goal in policy gradient is to find the best θ for a given policy πθ(s,a) with parameters θ. But how do we measure the quality of a given policy? There are 3 ways of measuring: <ul> <li>In episodic environments we can use the start value: J1(θ)=Vπθ(s1)=Eπθ[v1]</li> <li>In continuing environments we can use the average value: Javerage value(θ)=s∑dπθ(s)Vπθ(s)</li> <li>Or the average reward per time step Javerage reward(θ)=s∑dπθ(s)a∑πθ(s,a)Rsa</li> </ul> In the above, dπθ(s) is the stationary distribution of the markov chain for πθ. It tells us the amount of time we spend in each state, and so it provides the weighting required to get the average value or reward. <h2 id="policy-optimization"><a class="header" href="#policy-optimization">Policy Optimization</a></h2> Policy based reinforcement learning is an optimisation problem: <ul> <li>Gradient free methods: <ul> <li>Hill climbing</li> <li>Simplex / amoeba / Nelder Mead</li> <li>Genetic algorithms</li> </ul> </li> <li>Gradient methods are almost always more efficient: <ul> <li>Gradient descent</li> <li>Conjugate gradient</li> <li>Quasi Newton</li> </ul> </li> </ul> <h2 id="finite-difference-policy-gradient"><a class="header" href="#finite-difference-policy-gradient">Finite Difference Policy Gradient</a></h2> <ul> <li>Let J(θ) be any policy objective function</li> <li>Policy gradient algorithms search for a local maximum in J(θ) by ascending the gradient of the policy wrt parameters θ △θ=α∇θJ(θ)</li> <li>Where ∇θJ(θ) is the policy gradient (a vector of partial derivatives along each dimension) ∇θJ(θ)=<svg xmlns="http://www.w3.org/2000/svg" width='0.875em' height='4.800em' viewBox='0 0 875 4800'><path d='M863,9c0,-2,-2,-5,-6,-9c0,0,-17,0,-17,0c-12.7,0,-19.3,0.3,-20,1 c-5.3,5.3,-10.3,11,-15,17c-242.7,294.7,-395.3,682,-458,1162c-21.3,163.3,-33.3,349, -36,557 l0,1284c0.2,6,0,26,0,60c2,159.3,10,310.7,24,454c53.3,528,210, 949.7,470,1265c4.7,6,9.7,11.7,15,17c0.7,0.7,7,1,19,1c0,0,18,0,18,0c4,-4,6,-7,6,-9 c0,-2.7,-3.3,-8.7,-10,-18c-135.3,-192.7,-235.5,-414.3,-300.5,-665c-65,-250.7,-102.5, -544.7,-112.5,-882c-2,-104,-3,-167,-3,-189 l0,-1292c0,-162.7,5.7,-314,17,-454c20.7,-272,63.7,-513,129,-723c65.3, -210,155.3,-396.3,270,-559c6.7,-9.3,10,-15.3,10,-18z'/></svg>∂θ1∂J(θ)⋮∂θn∂J(θ)<svg xmlns="http://www.w3.org/2000/svg" width='0.875em' height='4.800em' viewBox='0 0 875 4800'><path d='M76,0c-16.7,0,-25,3,-25,9c0,2,2,6.3,6,13c21.3,28.7,42.3,60.3, 63,95c96.7,156.7,172.8,332.5,228.5,527.5c55.7,195,92.8,416.5,111.5,664.5 c11.3,139.3,17,290.7,17,454c0,28,1.7,43,3.3,45l0,1209 c-3,4,-3.3,16.7,-3.3,38c0,162,-5.7,313.7,-17,455c-18.7,248,-55.8,469.3,-111.5,664 c-55.7,194.7,-131.8,370.3,-228.5,527c-20.7,34.7,-41.7,66.3,-63,95c-2,3.3,-4,7,-6,11 c0,7.3,5.7,11,17,11c0,0,11,0,11,0c9.3,0,14.3,-0.3,15,-1c5.3,-5.3,10.3,-11,15,-17 c242.7,-294.7,395.3,-681.7,458,-1161c21.3,-164.7,33.3,-350.7,36,-558 l0,-1344c-2,-159.3,-10,-310.7,-24,-454c-53.3,-528,-210,-949.7, -470,-1265c-4.7,-6,-9.7,-11.7,-15,-17c-0.7,-0.7,-6.7,-1,-18,-1z'/></svg></li> </ul> The simplest way to compute the policy gradient is to use finite differences: <ul> <li>For each dimension k∈[1,n]: <ul> <li>Estimate the kth partial derivative of objective function wrt θ</li> <li>By perturbing θ by a small amount ϵ in the kth dimension ∂θk∂J(θ)≈J(θ+ϵuk)−J(θ)</li> <li>Where uk is a unit vector with 1 in the kth component and 0 elsewhere</li> </ul> </li> <li>This is not the most efficient algorithm, as it requires n evaluations (once for each dimension) to compute a single gradient step</li> <li>It is simple, noisy and inefficient, but sometimes works</li> <li>Works for arbitrary policies, even if the policy is not differentiable</li> </ul> <h2 id="likelihood-ratios"><a class="header" href="#likelihood-ratios">Likelihood Ratios</a></h2> We now want to compute the policy gradient analytically, assuming: <ul> <li>The policy πθ is differentiable whenever it is non-zero; and</li> <li>We know the gradient ∇θπθ(s,a)</li> <li>Likelihood ratio methods exploit the following identity (call it the log trick): ∇θπθ(s,a)=πθ(s,a)πθ(s,a)∇θπθ(s,a)=πθ(s,a)∇θlogπθ(s,a)</li> <li>Note that we use the simple identity ∂θlogf(θ)=f(θ)∂θf(θ)</li> <li>The new formulation is nicer to work with because we have πθ(s,a) on the left, which when integrated over, basically gives us the expectation over our policy πθ <ul> <li>This allows us to basically sample trajectories from the data and compute the gradient at each step</li> </ul> </li> <li>The score function is the quantity ∇θlogπθ(s,a)</li> </ul> <h2 id="linear-softmax-policy"><a class="header" href="#linear-softmax-policy">Linear Softmax Policy</a></h2> Use the softmax policy as a simple running example: <ul> <li>Weight actions using linear combination of features ϕ(s,a)⊺θ</li> <li>The probability of action is then proportional to the exponentiated weight: πθ(s,a)∝eϕ(s,a)⊺θ</li> <li>The score function is then: ∇θlogπθ(s,a)=ϕ(s,a)−Eπθ[ϕ(s,⋅)]</li> <li>Note that we are omitting the derivation of the second term of the score function which is a bit more involved, as it involves differentiating the normalization factor (not shown above)</li> </ul> <h2 id="gaussian-policy"><a class="header" href="#gaussian-policy">Gaussian Policy</a></h2> In continuous action spaces, a Gaussian policy is natural <ul> <li>Let the mean of the gaussian be a linear combination of state features μ(s)=ϕ(s)⊺θ</li> <li>The variance may be fixed σ2 or parametrized</li> <li>The policy is gaussian (recall that we are in a continuous action space, so a is a vector of floats): a∼N(μ(s),σ2)</li> <li>The score function is then: ∇θlogπθ(s,a)=σ2(a−μ(s))ϕ(s)</li> <li>We can probably derive this score function by writing down the PDF of the gaussian distribution for πθ(s,a) and then taking the derivative</li> </ul> <h2 id="policy-gradient-theorem-one-step-mdps"><a class="header" href="#policy-gradient-theorem-one-step-mdps">Policy Gradient Theorem: One-Step MDPs</a></h2> Consider a simple class of one-step MDPs to simplify the math <ul> <li>Start in a state s∼d(s)</li> <li>Terminate after one step with reward r=Rs,a</li> <li>This is a sort of contextual bandit</li> </ul> Use likelihood ratios to compute the policy gradient <ul> <li>First we pick our objective function, which is just the expected reward (averaged over our start state and action that we choose) J(θ)=Eπθ[r]=s∈S∑d(s)a∈A∑πθ(s,a)Rs,a</li> <li>Then we take the derivative to do gradient ascent: ∇θJ(θ)=s∈S∑d(s)a∈A∑∇θπθ(s,a)Rs,a=s∈S∑d(s)a∈A∑πθ(s,a)∇θlogπθ(s,a)Rs,a=Eπθ[∇θlogπθ(s,a)r]</li> <li>Note that when taking the gradient of πθ, we use the log trick to rewrite it in line 2, and it becomes a new expectation again because we recover πθ(s,a) outside of the gradient. This shows the power of the log trick.</li> </ul> <h2 id="policy-gradient-theorem"><a class="header" href="#policy-gradient-theorem">Policy Gradient Theorem</a></h2> But we don't just want to do one-step MDPs, we want to generalize to multi-step MDPs <ul> <li>It turns out that we just need to replace the instantaneous reward r with long term value Qπ(s,a) (I suppose this means we need to model Q as well)</li> <li>Regardless of whether we use the (i) start state objective, (ii) average value objective or (iii) average reward objective, the policy gradient theorem hold</li> </ul> <blockquote> Theorem. Policy Gradient Theorem. For any differentiable policy πθ(s,a), and for any of the policy objective functions J1, Javerage value or Javerage reward, the policy gradient is: ∇θJ(θ)=Eπθ[∇θlogπθ(s,a)Qπθ(s,a)] </blockquote> <h2 id="monte-carlo-policy-gradient-reinforce"><a class="header" href="#monte-carlo-policy-gradient-reinforce">Monte Carlo Policy Gradient (REINFORCE)</a></h2> The policy gradient theorem basically gives rise to a simple monte carlo policy gradient algorithm to find the optimal policy: <blockquote> REINFORCE Algorithm. <ul> <li>Initialize θ randomly</li> <li>For each episode {s1,a1,r2,...,sT−1,aT1,rT}∼πθ do: <ul> <li>For t=1 to T−1 do: <ul> <li>θ←θ+α∇θlogπθ(st,at)Gt</li> </ul> </li> </ul> </li> <li>Return θ</li> </ul> </blockquote> Note that: <ul> <li>We are doing monte carlo, i.e. we wait until the end of the episode before we go back to update the parameters for each time step.</li> <li>We are doing SGD, so there is no expectation term</li> <li>We use the return Gt as an unbiased sample of Qπθ(st,at). Recall that Gt is the total discounted reward from time step t until termination.</li> <li>This is the simplest and oldest policy gradient algorithm.</li> </ul> Empirically, policy gradient methods have a nice learning curve without the jittery behaviour of value based methods. But, monte carlo methods take very very long (millions of steps) to converge due to high variance. <h2 id="actor-critic-policy-gradient"><a class="header" href="#actor-critic-policy-gradient">Actor Critic Policy Gradient</a></h2> The main problem with monte carlo policy gradient is the high variance of the return Gt. Sometimes we get no reward, sometimes we get high reward. The idea is thus to use a critic to estimate the action-value function Q: Qw(s,a)≈Qπθ(s,a) The name critic refers to the value function, which simply "watches" and evaluates the value of an action, whilst the actor is the policy itself which decides how we should act. We maintain two sets of parameters: <ul> <li>Critic updates the action value function parameters w</li> <li>Actor updates the policy parameters θ, in the direction suggested by the critic</li> </ul> Actor-critic algorithms follow an approximate policy gradient: ∇θJ(θ)△θ≈Eπθ[∇θlogπθ(s,a)Qw(s,a)]=α∇θlogπθ(s,a)Qw(s,a) Notice that we just replace Gt with Qw(s,a), which is the value function of the critic model. <h2 id="estimating-the-action-value-function-of-the-critic"><a class="header" href="#estimating-the-action-value-function-of-the-critic">Estimating the Action-Value Function of the Critic</a></h2> The critic is solving a familiar problem: policy evaluation. What is the value for policy πθ for current parameters θ? We have explored this previously, i.e.: <ul> <li>Monte Carlo policy evaluation</li> <li>TD-λ</li> <li>or least squares policy evaluation</li> </ul> This leads us to a simple actor-critic algorithm: <ul> <li>Critic: use linear value function approximation Qw(s,a)=ϕ(s,a)⊺w <ul> <li>Update w using linear TD(0)</li> </ul> </li> <li>Actor: update θ using policy gradient</li> </ul> <blockquote> Q Actor Critic (QAC) Algorithm. <ul> <li>Initialize s,θ</li> <li>Sample a∼πθ</li> <li>for each step: <ul> <li>Sample reward r=Rsa; sample transitions s′∼Psa</li> <li>Sample action a′∼πθ(s′,a′)</li> <li>δ=r+γQw(s′,a′)−Qw(s,a)</li> <li>θ=θ+α∇θlogπθ(s,a)Qw(s,a)</li> <li>w←w+βδϕ(s,a)</li> <li>a←a′,s←s′</li> </ul> </li> </ul> </blockquote> Note that the δ line is simply the TD(0) update, and the θ line is simply the policy gradient step. <h2 id="bias-in-actor-critic-algorithms"><a class="header" href="#bias-in-actor-critic-algorithms">Bias in Actor Critic Algorithms</a></h2> The problem is that approximating the policy gradient introduces bias <ul> <li>A biased policy gradient may not find the right solution</li> <li>Luckily, if we choose the value function approximation carefully, we can avoid introducing any bias</li> <li>i.e. we can still follow the exact policy gradient (see Compatible Function Approximation theorem)</li> </ul> <h2 id="reducing-variance-using-a-baseline"><a class="header" href="#reducing-variance-using-a-baseline">Reducing Variance using a Baseline</a></h2> Here we move on to tricks to improve the training algorithm. The baseline method is the best known trick. The main idea is to subtract a baseline function B(s) from the policy gradient, such that we can reduce variance without changing expectation. The reason is that since B(s) does not depend on the action a, its expectation when plugged in will be 0, like so: Eπθ[∇θlogπθ(s,a)B(s)]=s∈S∑dπθ(s)a∈A∑∇θπθ(s,a)B(s)=s∈S∑dπθB(s)∇θa∈A∑πθ(s,a)=0 Note that in line 2, since πθ is a probability, the right-most term sums to 1, which is a constant. Hence the gradient ∇θ will become 0, and the expectation resolves to 0. A good and popular choice of baseline is the state value function B(s)=Vπθ(s). <ul> <li>So we can rewrite the policy gradient using the advantage function Aπθ(s,a): Aπθ(s,a)∇θJ(θ)=Qπθ(s,a)−Vπθ(s)=Eπθ[∇θlogπθ(s,a)Aπθ(s,a)]</li> <li>Intuitively, the advantage function tells us the additional benefit of an action over the baseline value of being in that state</li> </ul> <h2 id="estimating-the-advantage-function"><a class="header" href="#estimating-the-advantage-function">Estimating the Advantage Function</a></h2> The advantage function can significantly reduce the variance of the policy gradient. The naive way is to estimate Vπθ(s) and Qπθ(s,a) separately, i.e.: V^(s)Qw(s,a)A(s,a)≈Vπθ(s)≈Qπθ(s,a)=Qw(s,a)−V^(s) Then, we use TD methods to update both value functions. However, this is not efficient in terms of parameters. The better way is to observe that the TD error δπθ is an unbiased estimate of the advantage function. Therefore, we can plug in the TD error in our policy gradient update instead. <ul> <li>Recall that the TD error is δπθ=r+γVπθ(s′)−Vπθ(s)</li> <li>And this is an unbiased estimate of the advantage function Eπθ[δπθ∣s,a]=Eπθ[r+γVπθ(s′)∣s,a]−Vπθ(s)=Qπθ(s,a)−Vπθ(s)=Aπθ(s,a)</li> <li>Note that we have not taken any approximations above. We just showed that the expectation of the TD(0) error when following policy πθ corresponds to the advantage function. Interestingly, computing the TD error does not require estimating the Q function, only the V function. This gives us a simpler update</li> <li>We can thus use the TD error to compute the policy gradient ∇θJ(θ)=Eπθ[∇θlogπθ(s,a)δπθ]</li> <li>In practice, since we do not have the true value function Vπθ, we use an estimate V^: δw=r+γVw^(s′)−Vw^(s)</li> <li>This approach only requires one set of critic parameters w</li> </ul> <h2 id="improving-the-actor-updates"><a class="header" href="#improving-the-actor-updates">Improving the Actor Updates</a></h2> Recall that we can estimate the value function Vθ(s) from targets at different time scales to trade-off bias and variance: <ul> <li>For MC, the target is the return Gt △θ=α(Gt−Vθ(s))ϕ(s)</li> <li>For TD(0), the target is the TD target △θ=α(r+γVθ(s′)−Vθ(s))ϕ(s)</li> <li>For forward view TD(λ), the target is the lambda return △θ=α(vtλ−Vθ(s))ϕ(s)</li> <li>For backward view TD(λ), we use eligibility traces δtet△θ=rt+1+γV(st+1)−V(st)=λet−1+ϕ(st)=αδtet</li> </ul> Similarly, we can estimate the policy gradient for the actor at many time scales. The main target we want to estimate is ∇θJ(θ)=Eπθ[∇θlogπθ(s,a)Aπθ(s,a)] <ul> <li>MC policy gradient uses the error from the return △θ=α(Gt−V^(st))∇θlogπθ(st,at)</li> <li>Actor critic policy gradient uses the one-step TD error △θ=α(r+γV^(st+1)−V^(st))∇θlogπθ(st,at)</li> <li>Forward view TD(λ) uses the λ target △θ=α(vtλ−V^(st))∇θlogπθ(st,at)</li> <li>Backward view TD(λ) uses eligibility traces. Note that the eligbility update now uses the score function instead of ϕ which was the feature function δtet+1△θ=rt+1+γV^(st+1)−V^(st)=λet+∇θlogπθ(s,a)=αδtet</li> <li>The update can be performed online, unlike MC policy gradient</li> </ul> <h2 id="natural-policy-gradient"><a class="header" href="#natural-policy-gradient">Natural Policy Gradient</a></h2> The natural policy gradient is a parametrization independent approach. It finds ascent direction that is closest to vanilla gradient, when changing the policy by a small, fixed amount ∇θnatπθ(s,a)=Gθ−1∇θπθ(s,a) <ul> <li>Where Gθ is the fisher information matrix Gθ=Eπθ[∇θlogπθ(s,a)∇θlogπθ(s,a)⊺]</li> <li>Notice that this obviates the need for a critic, as Gθ is based on the actor itself</li> </ul> <h2 id="summary-1"><a class="header" href="#summary-1">Summary</a></h2> The policy gradient has many equivalent forms ∇θJ(θ)Gθ−1∇θJ(θ)=Eπθ[∇θlogπθ(s,a)Gt]=Eπθ[∇θlogπθ(s,a)Qw(s,a)]=Eπθ[∇θlogπθ(s,a)Aw(s,a)]=Eπθ[∇θlogπθ(s,a)δ]=Eπθ[∇θlogπθ(s,a)δe]=wREINFORCEQ Actor-CriticAdvantage Actor-CriticTD Actor-CriticTD(λ) Actor-CriticNatural Actor-Critic Each formulation leads to a different SGD algorithm. We can learn the critic using policy evaluation (e.g. MC or TD learning) to estimate Qπ,Aπ or Vπ. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="learning--planning"><a class="header" href="#learning--planning">Learning & Planning</a></h1> <a href="https://www.youtube.com/watch?v=ItMutbeOHtc&list=PLqYmG7hTraZDM-OYHWgPebj2MfCFzFObQ&index=8">Lecture 8: Integrating Learning and Planning</a> A more conceptual lecture. <ul> <li>Introduction</li> <li>Model based RL</li> <li>Integrated Architectures</li> <li>Simulation Based Search</li> </ul> <h2 id="model-based-rl"><a class="header" href="#model-based-rl">Model Based RL</a></h2> <ul> <li>Lecture 7: Learn policy directly from experience</li> <li>Lecture 4-6: Learn value function directly from experience</li> <li>This lecture: learn model directly from experience <ul> <li>The model describes an agent's understanding of the environment.</li> </ul> </li> <li>Use planning to construct a value function or policy</li> <li>Integrate learning and planning into a single architecture</li> </ul> Taxonomy <ul> <li>Model free RL <ul> <li>No model</li> <li>Learn value function / policy function from experience</li> </ul> </li> <li>Model-based RL <ul> <li>Learn a model from experience</li> <li>Plan value function / policy function from model. Use model to look ahead and think and get rewards.</li> </ul> </li> </ul> The learning diagram for Model-based RL is as follows: <pre class="mermaid">graph LR; A("value/policy") -- acting --> B("experience"); B("experience") -- model learning --> C("model"); C("model") -- planning --> A("value/policy"); </pre> Advantages of Model-Based RL <ul> <li>Can efficiently learn model by supervised learning methods <ul> <li>Consider a game like chess, where the model and game rule is quite simple</li> <li>But the value function is very hard, because moving one piece can dramatically change the value function</li> <li>So in some cases, the model is a more compact and useful representation of the problem</li> <li>The model simply predicts our next state, given the previous state and action. We can have a teacher which is the environment or game engine.</li> </ul> </li> <li>Can reason about model uncertainty</li> </ul> Disadvantage of Model-based RL: <ul> <li>First learn a model, then construct a value function, which gives us two sources of approximation error</li> </ul> <h2 id="what-is-a-model"><a class="header" href="#what-is-a-model">What is a Model?</a></h2> A model M is a representation of an MDP ⟨S,A,P,R⟩ parametrized by η. <ul> <li>Assume state space S and action space A are known</li> <li>So a model M=⟨Pη,Rη⟩ represents state transitions Pη≈P and rewards Rη≈R</li> <li>And we can use the model to sample next state and reward St+1Rt+1∼Pη(St+1∣St,At)∼Rη(Rt+1∣St,At)</li> <li>Note that we just need to advance one time step to start learning</li> <li>It is typical to assume conditional independence between state transitions and rewards, i.e. P[St+1,Rt+1∣St,At]=P[St+1∣St,At]P[Rt+1∣St,At]</li> </ul> <h2 id="model-learning"><a class="header" href="#model-learning">Model Learning</a></h2> The goal is to estimate the model Mη from experience {S1,A1,R2,...,ST} <ul> <li>Observe that this is just a supervised learning problem S1,A1S2,A2⋮ST−1,AT−1→R2,S2→R3,S3→RT,ST</li> <li>Learning s,a→r is a regression problem</li> <li>Learning s,a→s′ is a density estimation problem</li> <li>Pick our favourite loss function, e.g. MSE, KL divergence</li> <li>Pick parameters η that minimize the empirical loss</li> </ul> Examples of models: <ul> <li>Table lookup model</li> <li>Linear expectation model</li> <li>Linear gaussian model</li> <li>Gaussian process model</li> <li>Deep belief network model</li> </ul> <h2 id="table-lookup-model"><a class="header" href="#table-lookup-model">Table Lookup Model</a></h2> Our model is an explicit MDP, i.e. P^,R^ <ul> <li>We simply count the number of visits to each state action pair, i.e. N(s,a) and record the probability of each resulting s′ and mean reward Rt: P^s,s′aR^sa=N(s,a)1t=1∑T1(St=s,At=a,St+1=s′)=N(s,a)1t=1∑T1(St=s,At=a)Rt</li> <li>An alternative to do this in a non-parametric way: <ul> <li>At each time step t, record experience tuple ⟨St,At,Rt+1,St+1⟩</li> <li>To sample model, each time we are in state s,a, we randomly pick a tuple matching ⟨s,a,⋅,⋅⟩</li> </ul> </li> </ul> <h2 id="ab-example-1"><a class="header" href="#ab-example-1">AB Example</a></h2> Suppose we have 8 episodes of experience: <ul> <li><code>A, 0, B, 0</code></li> <li><code>B, 1</code></li> <li><code>B, 1</code></li> <li><code>B, 1</code></li> <li><code>B, 1</code></li> <li><code>B, 1</code></li> <li><code>B, 1</code></li> <li><code>B, 0</code></li> </ul> From this data, we learn a model using the one-step supervised learning method like so: <pre class="mermaid">graph LR; A("A") -- r=0, 100% --> B("B"); B("B") -- r=1, 75% --> C("END"); B("B") -- r=0, 25% --> D("END"); </pre> <h2 id="planning-with-a-model"><a class="header" href="#planning-with-a-model">Planning with a Model</a></h2> Given a model Mη=⟨Pη,Rη⟩, now we want to solve the MDP ⟨S,A,Pη,Rη⟩. We can use our favourite planning algo: <ul> <li>Value iteration</li> <li>Policy iteration</li> <li>Tree search</li> </ul> One of the simplest approaches is to do sample-based planning, but also one of the most powerful. The idea is to use the model only to generate samples. We sample experiences from the model: St+1Rt+1∼Pη(St+1∣St,At)∼Rη(Rt+1∣St,At) We can then apply model free RL to samples, e.g. using <ul> <li>Monte carlo control</li> <li>Sarsa</li> <li>Q-learning etc.</li> </ul> Sample based planning methods are often more efficient. <ul> <li>Sampling from the model is efficient because sampling gives us high probability events, compared to full width look up</li> </ul> <h2 id="back-to-ab-example"><a class="header" href="#back-to-ab-example">Back To AB Example</a></h2> After we built the model, we can sample from it, for e.g. <ul> <li><code>B, 1</code></li> <li><code>B, 0</code></li> <li><code>B, 1</code></li> <li><code>A, 0, B, 1</code></li> <li><code>B, 1</code></li> <li><code>A, 0, B, 1</code></li> <li><code>B, 1</code></li> <li><code>B, 0</code></li> </ul> Based on this sample data, we apply monte carlo learning, and end up with: V(A)=1,V(B)=0.75 <h2 id="planning-with-an-inaccurate-model"><a class="header" href="#planning-with-an-inaccurate-model">Planning with an Inaccurate Model</a></h2> Given an imperfect model ⟨Pη,Rη⟩=⟨P,R⟩, we need to remember that the performance of model-based RL is limited to the optimal policy for the approximate MDP ⟨S,A,Pη,Rη⟩. <ul> <li>i.e. the model based RL is only as good as the estimated model</li> <li>When the model is inaccurate, the planning process will compute a suboptimal policy</li> </ul> Solutions: <ol> <li>When model is wrong, use model-free RL</li> <li>Reason explicitly about model uncertainty</li> </ol> <h2 id="integrated-architectures"><a class="header" href="#integrated-architectures">Integrated Architectures</a></h2> Bring together the best of model-based and model free architectures. We consider two sources of experience: <ul> <li>Real experience. Sampled from environment (true MDP): S′R∼Ps,s′a=Rsa</li> <li>Simulated experience. Sampled from model (approximate MDP): S′R∼Pη(S′∣S,A)=Rη(R∣S,A)</li> </ul> Some taxonomy: <ul> <li>Model Free RL: <ul> <li>No model</li> <li>Learn value function (or policy) from real experience</li> </ul> </li> <li>Model based RL <ul> <li>Learn a model from real experience</li> <li>Plan value function (or policy) from simulated experience</li> </ul> </li> <li>Dyna <ul> <li>Learn a model from real experience</li> <li>Learn and plan value function (or policy) from real and simulated experience</li> </ul> </li> </ul> The dyna architecture looks like that: <pre class="mermaid">graph LR; A("value/policy") -- acting --> B("experience"); B("experience") -- model learning --> C("model"); C("model") -- planning --> A("value/policy"); B("experience") -- direct RL --> A("value/policy"); </pre> <blockquote> Dyna-Q Algorithm <ul> <li>Initialize Q(s,a) and Model(s,a) for all s∈S and a∈A(s)</li> <li>Do forever: <ol> <li>S←current state</li> <li>A←ϵ-greedy(S,Q)</li> <li>Execute action A; observe reward R and state S′</li> <li>Q(S,A)←Q(S,A)+α[R+γmaxaQ(S′,a)−Q(S,A)]</li> <li>Model(S,A)←R,S′</li> <li>Repeat n times: <ul> <li>S←random previously observed state</li> <li>A←random action previously taken in S</li> <li>R,S′←Model(S,A)</li> <li>Q(S,A)←Q(S,A)+α[R+γmaxaQ(S′,a)−Q(S,A)]</li> </ul> </li> </ol> </li> </ul> </blockquote> Note that: <ul> <li>Step 4 is the standard Q-learning update step</li> <li>Step 5 updates the model using simple SGD supervised learning</li> <li>Step 6 is the thinking/planning step, where we "imagine" scenarios using our model and update Q n times, without actually moving our agent in the real world</li> </ul> Experiments show that planning significantly speeds up convergence, requiring much fewer exploration steps in the real world to converge. So we are squeezing much more information out of what we have explored so far. A variation to Dyna-Q is Dyna-Q+, which puts higher weight on unexplored states and encourages exploration. <h2 id="simulation-based-search"><a class="header" href="#simulation-based-search">Simulation Based Search</a></h2> Two key ideas: sampling and forward search. Forward search algorithms select the best action by lookahead. A search tree is built with the current state st at the root. We then use a model of the MDP to look ahead. So we do not need to solve the entire MDP, just the sub-MDP starting from the current position. Solving the entire MDP is a waste of time. Note that this is in contrast to the Dyna-Q algorithm, where the "thinking" step starts by randomly visiting a previously observed state. So we can simulate episodes of experience from now using our model, and apply model free RL to simulated episodes. Simulation based search: <ul> <li>Simulate episodes of experience from now with the model {stk,Atk,Rt+1k,...,STk}k=1K∼Mv</li> <li>Apply model free RL to simulated episodes <ul> <li>If we use Monte-carlo control, we get Monte-Carlo search</li> <li>If we use Sarsa for control, we get TD search</li> </ul> </li> </ul> <h3 id="simple-monte-carlo-search"><a class="header" href="#simple-monte-carlo-search">Simple Monte Carlo Search</a></h3> <ul> <li>Given a model Mv and some simulation policy π (how we pick actions in our imagination)</li> <li>For each current action a∈A: <ul> <li>Simulate K episodes using model from current real state st {st,at,Rt+1k,St+1k,At+1k,...,STk}k=1K∼Mv,π</li> <li>Note that after current at, we follow policy π for future actions</li> <li>Evaluate actions by mean return (Monte Carlo evaluation) Q(st,a)=K1k=1∑KGt→qπ(st,a)</li> </ul> </li> <li>Select current real actions with maximum value at=a∈AargmaxQ(st,a)</li> </ul> <h3 id="monte-carlo-tree-search-evaluation"><a class="header" href="#monte-carlo-tree-search-evaluation">Monte Carlo Tree Search (Evaluation)</a></h3> MCTS differs in that we allow the policy to improve within our simulation (i.e. policy π is not stationary within our simulation runs). <ul> <li>Given a model Mv</li> <li>Simulate K episodes using model from current real state st using current simulation policy π {st,at,Rt+1k,St+1k,At+1k,...,STk}k=1K∼Mv,π</li> <li>Build a search tree containing visited states and actions from the above simulations</li> <li>Evaluate states Q(s,a) by mean return of episodes starting from s,a (i.e. monte carlo evaluation) Q(s,a)=N(s,a)1k=1∑Ku=t∑T1(Su=s,Au=a)Gu→qπ(s,a)</li> <li>Note that: <ul> <li>N(s,a) is the number of times we visited the s,a pair during our simulations.</li> <li>We are assuming our model is good enough so that we can rely on the simulated returns</li> <li>We are not really learning a persistent Q function. At each step, we run simulations to get fresh estimations of the Q values</li> </ul> </li> <li>After the search is finished, select the current real action with maximum value in the search tree at=a∈AargmaxQ(st,a)</li> </ul> In MCTS, the simulation policy π actually improves. <ul> <li>Each simulation comprises two phases (in-tree, out of tree) <ul> <li>Tree policy (improves): pick actions to maximize Q(S,A) by looking at the search tree and the node's children</li> <li>Default policy (fixed): in our simulation, when we run beyond the frontier of the search tree, we will pick actions randomly</li> </ul> </li> <li>Repeat (each simulation) <ul> <li>Evaluate states Q(S,A) by monte carlo evaluation</li> <li>Improve tree policy, e.g. by ϵ-greedy(Q)</li> </ul> </li> <li>Essentially, this is monte carlo control applied to simulated experience</li> <li>This method converges on the optimal search tree, i.e. Q(S,A)→q∗(S,A)</li> </ul> <h2 id="case-study-game-of-go"><a class="header" href="#case-study-game-of-go">Case Study: Game of Go</a></h2> Position evaluation in Go: <ul> <li>How good is a position s?</li> <li>Reward function (undiscounted): RtRT=0 for all non-terminal steps t<T={1 if black wins0 if white wins</li> <li>Policy π=⟨πB,πW⟩ selects moves for both players</li> <li>The value function (how good is position s): vπ(s)v∗(s)=Eπ[RT∣S=s]=P[Black wins∣S=s]=πBmaxπWminvπ(s)</li> </ul> How does simple monte carlo evaluation work? <ul> <li>Suppose we start with a certain board configuration s</li> <li>We simulate many runs of games starting from s with current policy π</li> <li>The value function vπ(s) would be the fraction of simulated games where black wins</li> </ul> <h2 id="how-then-does-monte-carlo-tree-search-work"><a class="header" href="#how-then-does-monte-carlo-tree-search-work">How then does monte carlo tree search work?</a></h2> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="pmpp"><a class="header" href="#pmpp">PMPP</a></h1> My notes on: <blockquote> Programming Massively Parallel Processors, 4th Edition by Wen-mei W. Hwu, David B. Kirk, Izzat El Hajj </blockquote> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="introduction-2"><a class="header" href="#introduction-2">Introduction</a></h1> Historically, vast majority of software applications are written as sequential programs. From 1980 to 2000, advancements in CPUs have driven performance increases in this paradigm. Since 2003, this has slowed down due to increased energy consumption and heat dissipation issues. Hence, there is a move toward multiple physical CPUs, referred to as processor cores, used in each chip. To benefit from multiple processor cores, there must be multiple instruction sequences that can simultaneously execute on these processor cores. Hence, application software that will continue to enjoy significant performance improvement will be parallel programs, in which multiple threads of execution cooperate to complete the work faster. There are two main trajectories for designing microprocessors: <ul> <li>Multicore trajectory seeks to maintain execution speed of sequential programs.</li> <li>Many-thread trajectory seeks to maximize execution throughput of parallel programs. The idea is to have large number of threads (e.g. NVIDIA Tesla A100 has tens of thousands of threads)</li> </ul> The many-thread paradigm leads the race of floating-point calculation throughput. As of 2016, the ratio of peak throughput between many-thread GPUs and multicore CPUs is about <code>10</code>. This performance gap motivates many applications to move parts of their software to GPU for execution. Fundamentally, the design of CPU and GPU are quite different. <ul> <li>CPUs are optimized for sequential processing, and much space on the chip is dedicated for the following functions <ul> <li>They make use of sophisticated control logic to allow instructions from a single thread to execute in parallel while maintaining the appearance of sequential execution</li> <li>Large cache memories are provided to reduce the instruction and data access latencies</li> </ul> </li> <li>GPUs are optimized for throughput. In general, reducing latency is more expensive than reducing throughput. Hence, the solution is to increase throughput by including a massive number of threads. <ul> <li>Pipelined memory channels and arithmetic operations are allowed to have longer latency, allowing for smaller chip area for memory access hardware and arithmetic units</li> <li>This allows us to have a much higher degree of parallelism compared to CPUs</li> </ul> </li> </ul> Given the relative benefits of each, one should expect that applications should use both CPUs and GPUs. This is why CUDA is designed to support joint CPU - GPU execution of an application. <h1 id="architecture-of-a-gpu"><a class="header" href="#architecture-of-a-gpu">Architecture of a GPU</a></h1> A GPU is organized as an array of highly threaded streaming multiprocessors (SMs). <ul> <li>Each SM has a number of streaming processors (SPs) that share control logic and instruction cache</li> <li>Each GPU comes with high bandwidth (and also higher latency) global memory which we will refer to as DRAM. In NVIDIA's recent Pascal architecture, this is refered to as High Bandwidth Memory (HBM).</li> </ul> There is also a CPU - GPU / GPU - GPU connection to allow data transfers. In the Pascal family, this is called NVLINK, which allows transfers of up to 40 GB/s per channel. As the size of GPU memory grows, applications increasingly keep their data in the global memory instead of relying on the CPU-GPU connection. <h1 id="speeding-up-applications"><a class="header" href="#speeding-up-applications">Speeding up Applications</a></h1> How much parallelism can speed up an application depends on the portion of the application that can be parallelized. If the time taken for a sequential program is T and we can only parallelize <code>30%</code> of it, then a 10x speed-up of the parallel portion will only result in <code>27%</code> execution time reduction: 0.7T+0.03T=0.73T The fact that the level of speedup one can achieve through parallel execution can be severely limited by the parallelizable portion is known as Amdahl's Law. <h1 id="goals-of-the-book"><a class="header" href="#goals-of-the-book">Goals of the book</a></h1> Much of the book is dedicated to techniques for developing high performance parallel code. Focus on computational thinking techniques that will be amenable to parallel computing. Second goal is to write parallel programs with correct functionality and reliability, and in a way that you can debug and inspect the code. CUDA encourages use of simple forms of barrier synchronization, memory consistency, and atomicity for managing parallelism. Third goal is scalability across future hardware designs. The key to such scaling is to regularize and localize memory data accesses to minimize consumption of critical resources and conflicts in updating data parameters. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="data-parallel-computing"><a class="header" href="#data-parallel-computing">Data Parallel Computing</a></h1> When applications run slowly, the issue is usually that there is too much data to process. In many cases, this data can be split up and dealt with independently. Such independent evaluation of different pieces of data is called data parallelism. <ul> <li>This is different from Task parallelism, which is the task decomposition of applications. For example, we may need to do a vector multiplication and addition on the same set of data. In general, data parallelism is usually the main source of scalability for parallel programs.</li> </ul> Let us use a running example using color to gray scale conversion. This is a simple data parallel task as each pixel can be dealt with independently. Suppose we have a color image comprising many pixels, each pixel has an <code>r, g, b</code> value ranging from <code>0</code> to <code>1</code> indicating the intensity. We use a simple grayscale formula (note that the grayscale formula is because our eyes perceive luminosity of different colors differently): L=0.21r+0.72g+0.07b <h2 id="cuda-c-program"><a class="header" href="#cuda-c-program">CUDA C Program</a></h2> CUDA C is an extension of ANSI C programming language with minimal new syntax. The structure of a CUDA C program reflects the coexistence of a host (CPU) and one or more devices (GPUs) in the computer. Each source file can have a mixture of host code and device code. By definition, a traditional C program is a CUDA program that contains only host code. The device code is clearly marked with CUDA C keywords. When a kernel function is called, a large number of threads are launched on a device to execute the kernel. All the threads that are launched by a kernel call are collectively called a grid. When all threads of a grid have completed, the grid terminates and the execution continues on the host until the next grid is launched. Launching a grid generates many threads. In the grayscale example, each thread processes one pixel. One can assume that CUDA threads take very few clock cycles to generate and schedule. In contrast, traditional C programs can take thousands of clock cycles to generate and schedule. <blockquote> Threads. A thread is a simplified view of how a processor executes a sequential program. A thread consists of the code of the program, the point in the code that is being executed, and the values of its variables and data structures. The execution of a thread may be viewed as sequential. </blockquote> <h2 id="vector-addition"><a class="header" href="#vector-addition">Vector Addition</a></h2> Vector addition is the hello world example. Here is a simple vector addition in C: <blockquote> In CUDA programms, we suffix with <code>_h</code> to denote host variables and <code>_d</code> for device variables. </blockquote> <pre><code class="language-c">void vecAdd(float* A_h, float* B_h, float* C_h, int n) { for (int i = 0; i < n; ++1) { C_h[i] = A_h[i] + B_h[i]; } } int main() { // memory allocation etc. vecAdd(A, B, C, N); } </code></pre> In the vector addition example, we use a simple for loop to do addition. After execution, we can access the modified contents of <code>C</code> to see the results. A straightforward way to do vector addition in parallel is as follows: <ul> <li>Allocate device memory for A, B and C</li> <li>Copy A and B to device memory</li> <li>Call a kernel to launch threads to perform the actual addition</li> <li>Copy C from the device memory back to host</li> <li>Free device vector</li> </ul> <h2 id="device-global-memory"><a class="header" href="#device-global-memory">Device Global Memory</a></h2> Each device comes with their own dynamic random access memory called device global memory. e.g. NVIDIA Volta V100 comes with <code>16GB</code> or <code>32GB</code> of global memory. Before calling the addition kernel, we need to allocate space in the device global memory and transfer data from the host memory to the allocated space on the device. <ul> <li><code>cudaMalloc()</code> is used to allocate object on device global memory <ul> <li>Takes in address of a pointer to the allocated object</li> <li>Size of the allocated object in bytes</li> </ul> </li> <li><code>cudaFree()</code> is used to free object from global memory <ul> <li>Takes in the pointer to object to be freed</li> </ul> </li> </ul> <pre><code class="language-c">float *A_d // Pointer to a block of floating point numbers int size = n * sizeof(float); // n * 4 bytes cudaMalloc((void**) &A_d, size); ... cudaFree(A_d); </code></pre> Note: <ul> <li><code>(void**)</code> casts to a generic pointer, which is expected by <code>cudaMalloc</code> which frees any generic object</li> <li><code>&A_d</code> means the address of the pointer <code>A_d</code>: <ul> <li><code>A_d</code> itself is a pointer (on CPU memory) pointing to some spot on GPU memory</li> <li>We need to pass the address of <code>A_d</code> to <code>cudaMalloc</code> so that <code>cudaMalloc</code> can modify the contents of <code>A_d</code> to point at the newly allocated GPU memory location</li> <li>So <code>cudaMalloc</code> expects the address of a pointer</li> </ul> </li> <li><code>cudaFree</code> in contrast doesn't need to modify <code>A_d</code>, it just needs to go to the GPU location and free up the GPU memory there <ul> <li>Thus the argument to <code>cudaFree</code> is <code>A_d</code> instead of <code>&A_d</code></li> </ul> </li> <li>Note that <code>A_d</code> is a pointer to device memory, so it cannot be de-referenced in host code <ul> <li>Dereferencing using <code>*A_d</code> in host code will lead to errors</li> </ul> </li> </ul> After space has been allocated on device, we can transfer data from host to device. <ul> <li><code>cudaMemcpy()</code> <ul> <li>Pointer to destination</li> <li>Pointer to source</li> <li>Number of bytes copied</li> <li>Type / direction of transfer</li> </ul> </li> </ul> Now our <code>vecAdd</code> functio looks like this: <pre><code class="language-c">void vecAdd(float* A_h, float* B_h, float* C_h, int n) { int size = n * sizeof(float); float *A_d, *B_d, *C_d; cudaMalloc((void **) &A_d, size); cudaMalloc((void **) &B_d, size); cudaMalloc((void **) &C_d, size); cudaMemcpy(A_d, A_h, size, cudaMemcpyHostToDevice); cudaMemcpy(B_d, B_h, size, cudaMemcpyHostToDevice); // Invoke kernel cudaMemcpy(C_h, C_d, size, cudaMemcpyDeviceToHost); cudaFree(A_d); cudaFree(B_d); cudaFree(C_d); } </code></pre> The above code is quite self explanatory. This kind of host code that delegates work to a device is called a stub. Note that <code>cudaMemcpyDeviceToHost</code> etc. are special CUDA keywords. <blockquote> Notice that C syntax is a bit confusing - we write <code>float *A_d, *B_d, *C_d</code> where the <code>*</code> is attached to the variable, but in function parameters we write <code>void vecAdd(float* A_h, ...)</code>. It turns out that both ways of writing compile to the same code, and in both cases we are referring to <code>A_h, A_d</code> as float pointers. The reason why <code>*</code> is bound to the variable when declaring is to avoid the subtle bug of <code>float* A_d, B_d</code>, which actually declares <code>B_d</code> as a float, not a float pointer! </blockquote> <h2 id="kernel-functions-and-threading"><a class="header" href="#kernel-functions-and-threading">Kernel Functions and Threading</a></h2> Since all threads execute the same code, CUDA C programming is an instance of Single Program, Multiple Data (SPMD) parallel programming style. CUDA launches a grid of threads that are organized into a two level hierarchy: <ul> <li>Each grid is organized as an array of thread blocks, which are also referred to as blocks</li> <li>All blocks are of the same size, and each block can contain up to <code>1,024</code> threads</li> </ul> The total number of threads per block is specified by the host code when a kernel is called. For a given grid of threads, the number of threads in a block is available in a variable called <code>blockDim</code>. <ul> <li><code>blockDim</code> is a <code>struct</code> with three unsigned integer fields <code>x, y, z</code>, which may be accessed like <code>blockDim.x</code></li> <li>The choice of dimensionality for organizing threads usually reflects the dimensionality of the data</li> <li>In general, the number of threads is recommended to be a multple of <code>32</code> for hardware reasons</li> </ul> CUDA kernels and thus CUDA kernel code have access to two more built-in variables called <code>threadIdx</code> and <code>blockIdx</code>: <ul> <li><code>threadIdx</code> gives each thread a unique coordinate within a block <ul> <li>The first thread in a block has value <code>0</code> in <code>threadIdx.x</code> variable, second thread has <code>1</code> and so on</li> </ul> </li> <li><code>blockIdx</code> gives all threads in a block a common block coordinate <ul> <li>All threads in the first block have value <code>0</code> in their <code>blockIdx.x</code> variable</li> <li>Second block has value <code>1</code> and so on</li> </ul> </li> <li>Thus we can combine <code>threadIdx</code> and <code>blockIdx</code> to give a unique global index to each thread</li> </ul> Now we are ready to write kernel code for vector addition: <pre><code class="language-c">__global__ void vecAddKernel(float* A, float* B, float* C, int n) { int i = threadIdx.x + blockDim.x * blockIdx.x; if (i < n) { C[i] = A[i] + B[i]; } } </code></pre> Notes: <ul> <li>We do not need suffix <code>_d</code> since there is no confusion - kernel code only runs on device</li> <li><code>i</code> is a global index with a different value for each thread</li> <li><code>blockDim.x</code> is <code>256</code> for our example, provided by the CUDA runtime <ul> <li>This is specified in host code, when we launch the kernel</li> </ul> </li> <li>The keyword <code>__global__</code> indicates that this function is a kernel <ul> <li>In general, there are three qualifier keywords in CUDA C</li> <li><code>__host__</code> means callable from host and executed on host (i.e. normal C function)</li> <li><code>__device__</code> means callable from device and executed on device</li> <li><code>__global__</code> means callable from host / device and executed on device</li> </ul> </li> <li>Note that <code>if (i < n)</code> is a guarding condition <ul> <li>When we launch a grid, the number of threads is a multiple of the block size</li> <li>The number of threads has to be greater than or equal to <code>n</code></li> <li>To avoid errors on the excess threads, the conditional statement is used</li> </ul> </li> </ul> Notice that the CUDA kernel function has no for-loop. The reason is that the loop has been replaced by the grid of threads. Each thread in the grid corresponds to one iteration of the for loop. This is sometimes referred to as loop parallelism. <h2 id="complete-version-of-vecadd"><a class="header" href="#complete-version-of-vecadd">Complete Version of <code>vecAdd</code></a></h2> Now we can write the complete version of <code>vecAdd</code>. <pre><code class="language-c">void vecAdd(float* A, float* B, float* C, int n) { int size = n * sizeof(float); float *A_d, *B_d, *C_d; cudaMalloc((void **) &A_d, size); cudaMalloc((void **) &B_d, size); cudaMalloc((void **) &C_d, size); cudaMemcpy(A_d, A, size, cudaMemcpyHostToDevice); cudaMemcpy(B_d, B, size, cudaMemcpyHostToDevice); vecAddKernel<<<ceil(n / 256.0), 256>>>(A_d, B_d, C_d, n); cudaMemcpy(C, C_d, size, cudaMemcpyDeviceToHost); cudaFree(A_d); cudaFree(B_d); cudaFree(C_d); } </code></pre> Note that we call the kernel function with <code>vecAddKernel</code> and specify execution configuration parameters using the "<<< >>>" demarcators. <ul> <li>The first parameter <code>ceil(n / 256.0)</code> gives the number of blocks in the grid <ul> <li>Using <code>ceil</code> ensures that we have <code>>= n</code> threads</li> <li>Note that <code>256.0</code> is a float so that the division gives a float</li> </ul> </li> <li>The second specifies the number of threads in a block</li> </ul> If we vary <code>n</code>, the number of thread blocks will vary. A small GPU may execute <code>2</code> thread blocks in parallel, larger GPUs may do <code>64</code> or <code>128</code> at a time. In reality, the sequential way of doing <code>vecAdd</code> will probably be faster than the CUDA way, due to the overheads of allocating memory, data transfer, deallocating memory etc. Real CUDA programs do a lot more computation within the kernel to make the overheads worth it. <h2 id="compilation"><a class="header" href="#compilation">Compilation</a></h2> A traditional C compiler will not work with CUDA code. We need to use NVCC (NVIDIA C Compiler). The host code portions are compiled with standard C/C++ compilers and run as traditional CPU process. The device code is compiled by NVCC into <code>.ptx</code> files which are further compiled by a runtime component of NVCC into real object files. <h2 id="exercises"><a class="header" href="#exercises">Exercises</a></h2> <ol> <li>If we want to use each thread in a grid to calculate one output element of a vector addition, what would be the expression for mapping the thread/block indices to the data index (i)?</li> </ol> <pre><code class="language-c">i = blockIdx.x * blockDim.x + threadIdx.x; </code></pre> <ol start="2"> <li>Assume that we want to use each thread to calculate two adjacent elements of a vector addition. What would be the expression for mapping the thread/block indices to the data index (i) of the first element to be processed by a thread?</li> </ol> <pre><code class="language-c">i = (blockIdx.x * blockDim.x + threadIdx.x) * 2; </code></pre> <ol start="3"> <li>We want to use each thread to calculate two elements of a vector addition. Each thread block processes 2*blockDim.x consecutive elements that form two sections. All threads in each block will process a section first, each processing one element. They will then all move to the next section, each Page 45processing one element. Assume that variable i should be the index for the first element to be processed by a thread. What would be the expression for mapping the thread/block indices to data index of the first element?</li> </ol> <pre><code class="language-c">i = blockIdx.x * blockDim.x * 2 + threadIdx.x; </code></pre> <ol start="4"> <li>For a vector addition, assume that the vector length is 8000, each thread calculates one output element, and the thread block size is 1024 threads. The programmer configures the kernel call to have a minimum number of thread blocks to cover all output elements. How many threads will be in the grid?</li> </ol> <pre><code class="language-c">8192 </code></pre> <ol start="5"> <li>If we want to allocate an array of v integer elements in the CUDA device global memory, what would be an appropriate expression for the second argument of the cudaMalloc call?</li> </ol> <pre><code class="language-c">v * sizeof(int) </code></pre> <ol start="6"> <li>If we want to allocate an array of n floating-point elements and have a floating-point pointer variable A_d to point to the allocated memory, what would be an appropriate expression for the first argument of the cudaMalloc() call?</li> </ol> <pre><code class="language-c">(void **) &A_d </code></pre> <ol start="7"> <li>If we want to copy 3000 bytes of data from host array A_h (A_h is a pointer to element 0 of the source array) to device array A_d (A_d is a pointer to element 0 of the destination array), what would be an appropriate API call for this data copy in CUDA?</li> </ol> <pre><code class="language-c">cudaMemcpy(A_d, A_h, 3000, cudaMemcpyHostToDevice) </code></pre> <ol start="8"> <li>How would one declare a variable err that can appropriately receive the returned value of a CUDA API call?</li> </ol> <pre><code class="language-c">cudaError_t err; </code></pre> <ol start="9"> <li>Consider the following CUDA kernel and the corresponding host function that calls it:</li> </ol> <pre><code class="language-c">01 __global__ void foo_kernel(float* a, float* b, unsigned int N){ 02 unsigned int i=blockIdx.x*blockDim.x + threadIdx.x; 03 if(i < N) { 04 b[i]=2.7f*a[i] - 4.3f; 05 } 06 } 07 void foo(float* a_d, float* b_d) { 08 unsigned int N=200000; 09 foo_kernel <<< (N + 128–1)/128, 128 >>>(a_d, b_d, N); 10 } </code></pre> a. What is the number of threads per block? <pre><code>128 </code></pre> b. What is the number of threads in the grid? <pre><code>1,563 * 128 = 200,064 </code></pre> c. What is the number of blocks in the grid? <pre><code>1,563 </code></pre> d. What is the number of threads that execute the code on line 02? <pre><code>200,064 </code></pre> e. What is the number of threads that execute the code on line 04? <pre><code>200,000 </code></pre> <ol start="10"> <li>A new summer intern was frustrated with CUDA. He has been complaining that CUDA is very tedious. He had to declare many functions that he plans to execute on both the host and the device twice, once as a host function and once as a device function. What is your response?</li> </ol> <pre><code>We can use dual qualifiers to declare a function that works on both host and device. e.g. __host__ __device__ void myFunction(...) { // shared logic usable on both host and device } </code></pre> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="03---multidimensional-grids"><a class="header" href="#03---multidimensional-grids">03 - Multidimensional grids</a></h1> In the previous chapter we considered a simple 1D kernel. In general, a grid is a three-dimensional array of blocks, and each block is a 3D array of threads. We can specify them like so: <pre><code class="language-c">dim3 dimGrid(32, 1, 1); dim3 dimBlock(128, 1, 1); vecAddKernel<<<dimGrid, dimBlock>>>(...) </code></pre> This specifies a grid with <code>32</code> blocks and each block has <code>128</code> threads. So the total number of threads in the grid is <code>4,096</code>. Or we can let the variables be a function of something else, e.g. if <code>n</code> denotes the size of the vector: <pre><code class="language-c">dim3 dimGrid(ceil(n/256.0), 1, 1); dim3 dimBlock(256, 1, 1); vecAddKernel<<<dimGrid, dimBlock>>>(...) </code></pre> This fixes the number of threads per block at <code>256</code> and the number of blocks varies with the size of the vector. For convenience, if we are specifying 1D grids and blocks, we can just use a single number to represent it. In CUDA C, the allowable range of values: <ul> <li><code>gridDim.x</code> varies from 1 to 231−1</li> <li><code>gridDim.y</code>, <code>gridDim.z</code> varies from 1 to 216−1=65,535</li> <li>Consequently, <code>blockDim.x</code> ranges from 0 to <code>gridDim.x-1</code> and so on</li> </ul> The total size of a block in current CUDA systems is limited to <code>1024</code> threads. These threads can be distributed across three dimensions in any way as long as the total number of threads does not exceed <code>1024</code>. The dimensions of the number of blocks in a grid and the dimensions of the number of threads in a block need not match up. Sometimes, we need to refer to the index of a specific block in the grid or a specific thread in a block. The idiomatic way is: <ul> <li>Block: <code>(blockIdx.z, blockIdx.y, blockIdx.x)</code></li> <li>Thread: <code>(threadIdx.z, threadIdx.y, threadIdx.z)</code></li> </ul> Notice that the order of the index is reverse to the order of specifying the <code>dimGrid</code> or <code>dimBlock</code>. There is some reason for this, although it is not clear to me right now. <h2 id="mapping-threads-to-multidimensional-data"><a class="header" href="#mapping-threads-to-multidimensional-data">Mapping Threads to Multidimensional Data</a></h2> The choice of 1D, 2D or 3D thread organization is usually based on the nature of the data. For example, pictures are 2D and so it is conventional to choose 2D arrangements. Suppose we have a picture with <code>62 x 76</code> pixels, i.e. <code>62</code> in the <code>y</code> dimension and <code>76</code> in the <code>x</code> dimension. Assume we use a <code>16 x 16</code> block, with <code>16</code> threads in the <code>x, y</code> dimensions respectively. We will thus need <code>4</code> blocks in the <code>y</code> direction, and <code>5</code> blocks in the <code>x</code> direction. So for each thread: <ul> <li>The vertical (row) coordinate is: <code>row = blockIdx.y * blockDim.y + threadIdx.y</code></li> <li>The horizontal (column) coordinate is: <code>col = blockIdx.x * blockDim.x + threadIdx.x</code></li> </ul> How does the kernel call look like? Suppose we have: <ul> <li><code>n</code> represents number of pixels in the <code>y</code> direction</li> <li><code>m</code> represents the number of pixels in the <code>x</code> direction</li> <li>The picture data has been copied to device memory stored at pointer <code>Pin_d</code></li> </ul> Then the host code looks like: <pre><code class="language-c">dim3 dimGrid(ceil(m/16.0), ceil(n/16.0), 1); dim3 dimBlock(16, 16, 1); colorToGrayScaleConversion<<<dimGrid, dimBlock>>>(Pin_d, Pout_d, m, n); </code></pre> <h2 id="memory-layout"><a class="header" href="#memory-layout">Memory Layout</a></h2> Given that we have computed the row and col coordinates for each thread, ideally we would access the elements of <code>Pin_d</code> like <code>Pin_d[row][col]</code>. Unfortunately C does not allow us to do this - we are only allowed to access it as a 2D array in this manner if the dimensions of <code>Pin_d</code> is known at compile time. <blockquote> Memory Space. A memory space is a simplified view of how a processor accesses its memory. A memory space is usually associated with each running application. Each location has a byte and an address. Variables that need multiple bytes are stored in consecutive byte locations. Since there is only one address for each location, we say that the memory space has a flat layout. Thus, all multidimensional arrays are ultimately flattened into equivalent one dimensional arrays. When a C programmer uses multidimensional syntax to access an element, the transpiler is really translating these accesses into a 1D base pointer. </blockquote> Hence, the CUDA programmer needs to explicitly linearize the 2D accesses into a 1D access. There are two ways to do this: <ul> <li>row-major layout: rows are placed one after another consecutively in memory space</li> <li>column-major layout: columns are placed one after another consecutively in memory space</li> </ul> CUDA C uses the row-major layout; an example of the column-major layout is FORTRAN. We will thus focus on the row major layout. Let us denote the element of interest as Mj,i, i.e. the element with <code>y</code> coordinate <code>j</code> and <code>x</code> coordinate <code>i</code>. Suppose M is <code>4 x 4</code> and linearized into <code>16</code> element 1D array. In the flattened row-major layout, we access Mj,i using: <pre><code>j * 4 + i </code></pre> Now we'll tackle the 2D version of <code>colorToGrayscaleConversion</code>. <pre><code class="language-c">__global__ void colorToGrayscaleConversion(unsigned char * Pout, unsigned char * Pin, int width, int height) { int col = blockIdx.x * blockDim.x + threadIdx.x; int row = blockIdx.y * blockDim.y + threadIdx.y; if (col < width && row < height) { int grayOffset = row*width + col; int rgbOffset = grayOffset * CHANNELS; unsigned char r = Pin[rgbOffset]; unsigned char g = Pin[rgbOffset + 1]; unsigned char b = Pin[rgbOffset + 2]; Pout[grayOffset] = 0.21f*r + 0.71f*g + 0.07f*b; } } </code></pre> Notes: <ul> <li>Note that the input image is encoded as unsigned chars, since they encode a value from <code>0-255</code> <ul> <li>Each pixel is represented as <code>3</code> consecutive chars representing the <code>3</code> channels</li> <li>Thus pixel <code>i</code> starts at the char located at <code>i*3</code> position</li> </ul> </li> <li>We first compute <code>grayOffset</code>, which corresponds to the location for the pixel we are processing</li> <li><code>rgbOffset</code> is <code>3x grayOffset</code> because each pixel has 3 unsigned chars in the input image</li> <li>After extracting the <code>r, g, b</code> values from the input image, we compute and store the gray intensity into <code>Pout</code> using the linearized 1D acess <code>grayOffset</code></li> </ul> <h2 id="3d-arrays"><a class="header" href="#3d-arrays">3D arrays</a></h2> We can easily extend the discussion of 2D arrays into 3D. The idea is that the <code>x * y</code> coordinates form a plane, and we insert each "plane" one after another into the address space. <ul> <li>This means that we track a new <pre><code class="language-c">int plane = blockIdx.z * blockDim.z + threadIdx.z </code></pre> </li> <li>And access the 1D index as <pre><code class="language-c">P[plane*m*n + row*m + col] </code></pre> </li> </ul> <h2 id="image-blur"><a class="header" href="#image-blur">Image Blur</a></h2> Image blurring is closer to a real world example. The idea of blurring is to smooth out abrupt variations in pixel values. Mathematically, it is just computing the value of an output pixel as a weightedsum of a patch of pixels surrounding the pixel in the input image. In our simple example, we will take the simple average value of the <code>N x N</code> patch of pixels surrounding our target pixel. Similar to before, each thread will handle one output pixel. Thus much of the code will resemble that in the grayscale example. <pre><code class="language-c">__global__ void blurKernel(unsigned char *in, unsigned char *out, int w, int h) { int col = blockIdx.x * blockDim.x + threadIdx.x; int row = blockIdx.y * blockDim.y + threadIdx.y; if (col < w && row < h) { int pixVal = 0; int pixels = 0; for (int blurRow=-BLUR_SIZE; blurRow<BLUR_SIZE+1; ++blurRow) { for (int blurCol=-BLUR_SIZE; blurCol<BLUR_SIZE+1; ++blurCol) { int curRow = row + blurRow; int curCol = col + blurCol; if (curRow>=0 && curRow<h && curCol>=0 && curCol<w) { pixVal += in[curRow*w + curCol]; ++pixels; } } } out[row*w + col] = (unsigned char) (pixVal/pixels); } } </code></pre> Code walkthrough <ul> <li><code>in</code> and <code>out</code> point to the input and output pixel representation of the same size</li> <li>The <code>row</code> and <code>col</code> indices are derived as per normal from <code>blockIdx</code> and <code>threadIdx</code></li> <li><code>pixVal</code> stores the accumulated intensity, <code>pixels</code> stores the number of valid pixels <ul> <li>e.g. near the edge, <code>pixels</code> might not represent the full square patch</li> </ul> </li> <li>At each pixel location <code>in[row, col]</code>, we create a nested for loop to look up a square patch around the pixel <ul> <li>e.g. <code>BLUR_SIZE=1</code> will walk through a <code>3x3</code> patch around the centre pixel</li> <li>Notice that instead of a 2D index, we need to linearize into row-major form, i.e. <code>in[curRow*w + curCol]</code></li> <li>We accumulate intensity values into <code>pixVal</code> and count pixels walked in <code>pixels</code></li> <li>Finally we assign the average intensity into the corresponding location in <code>out</code></li> </ul> </li> </ul> <h2 id="matrix-multiplication"><a class="header" href="#matrix-multiplication">Matrix Multiplication</a></h2> Matrix multiplication is an important component of Basic Linear Algebra Subprograms (BLAS). Matrix multiplication between an <code>I x j</code> matrix <code>M</code> and a <code>j x k</code> matrix <code>N</code> produces a <code>I x k</code> matrix <code>P</code>. For simplicity, let's consider the case where both <code>N</code> and <code>M</code> are square with <code>Width</code> dimensions. <pre><code class="language-c">__global__ void MatrixMulKernel (float* M, float* N, float* P, int Width) { int row = blockIdx.y*blockDim.y + threadIdx.y; int col = blockIdx.x*blockDim.x + threadIdx.x; if ((row < Width) && (col < Width)) { float Pvalue = 0; for (int k = 0; k < Width; ++k) { Pvalue += M[row*Width + k] * N[k*Width + col]; } P[row*Width + col] = Pvalue; } } </code></pre> Code walkthrough: <ul> <li><code>row</code> and <code>col</code> are derived as per usual <ul> <li>For one position in the output matrix <code>P[row, col]</code></li> <li>We dot product <code>M[row, :]</code> with <code>N[:, col]</code></li> </ul> </li> <li>The inner index <code>k</code> loops through the columns of <code>M</code> and the rows of <code>N</code> to compute the dot product</li> <li>Note the row-major linearized accesses to <code>M</code>, <code>N</code> and <code>P</code></li> </ul> Note that for a given block of threads, it will process some contiguous rows from <code>M</code> and some contiguous columns from <code>N</code>. This means that visually, each block will process a tile of elements in the output matrix <code>P</code>. <h2 id="exercises-1"><a class="header" href="#exercises-1">Exercises</a></h2> <blockquote> <ol> <li>In this chapter we implemented a matrix multiplication kernel that has each thread produce one output matrix element. In this question, you will implement different matrix-matrix multiplication kernels and compare them.</li> </ol> a. Write a kernel that has each thread produce one output matrix row. Fill in the execution configuration parameters for the design. </blockquote> Suppose we have input matrices <code>M</code> and <code>N</code> which are both square with width <code>Width</code>. Since each thread does a whole row, we just need to use one row of <code>M</code> per thread but work through the entire <code>N</code> matrix. The configuration parameters would be <code>MatrixMulKernelRowWise<<<ceil(Width / n_threads), n_threads>>></code>, since we just need one thread per row of <code>M</code>. <pre><code class="language-c">__global__ void MatrixMulKernelRowWise (float* M, float* N, float* P, int Width) { int row = blockIdx.x*blockDim.x + threadIdx.x; if (row < Width) { for (int col = 0; col < Width; ++col) { float Pvalue = 0; for (int k = 0; k < Width; ++k) { Pvalue += M[row*Width + k] * N[k*Width + col]; } P[row*Width + col] = Pvalue; } } } </code></pre> <blockquote> b. Write a kernel that has each thread produce one output matrix column. Fill in the execution configuration parameters for the design. </blockquote> This is very similar, except that we work through the whole <code>M</code> matrix and only use a single column from <code>N</code>. The configuration parameters would be <code>MatrixMulKernelColWise<<<ceil(Width / n_threads), n_threads>>></code>, since we just need one thread per col of <code>N</code>. <pre><code class="language-c">__global__ void MatrixMulKernelColWise (float* M, float* N, float* P, int Width) { int col = blockIdx.x*blockDim.x + threadIdx.x; if (col < Width) { for (int row = 0; row < Width; ++row) { float Pvalue = 0; for (int k = 0; k < Width; ++k) { Pvalue += M[row*Width + k] * N[k*Width + col]; } P[row*Width + col] = Pvalue; } } } </code></pre> <blockquote> c. Analyze the pros and cons of each of the two kernel designs. </blockquote> Not very sure about this, but my thinking is that the row-wise version is more efficient. Note that access to <code>N</code> is strided in both cases, so the difference is in how we access <code>M</code>. Since each thread processes one row of <code>M</code>, we get the following: <ul> <li>thread 1 accesses <code>M[0, 0]</code>. Increment <code>k</code> and we access <code>M[0, 1]</code>, which is the next element in memory layout.</li> <li>thread 2 accesses <code>M[1, 0]</code>. Increment <code>k</code> and we access <code>M[1, 1]</code>, which is next element in memory layout</li> </ul> <blockquote> </blockquote> So as we increment <code>k</code>, each thread can benefit from its local cache and read consecutive elements efficiently. In contrast, for col-wise: <ul> <li>We start at row <code>0</code>, and all threads access <code>M[0, 0]</code>. The access is efficient as we increment <code>k</code>.</li> <li>But for each subsequent row, we need to read a row of <code>M</code> from global memory again</li> <li>So we incur much more global memory accesses than the row-wise approach</li> </ul> <blockquote> <ol start="2"> <li>A matrix-vector multiplication takes an input matrix B and a vector C and produces one output vector A. Each element of the output vector A is the dot product of one row of the input matrix B and C, that is, A[i]=∑jB[i][j]+C[j]. For simplicity we will handle only square matrices whose elements are single-precision floating-point numbers. Write a matrix-vector multiplication kernel and the host stub function that can be called with four parameters: pointer to the output matrix, pointer to the input matrix, pointer to the input vector, and the number of elements in each dimension. Use one thread to calculate an output vector element.</li> </ol> </blockquote> Assuming that <code>B</code> has dimensions <code>h x w</code> and thus <code>C</code> has dimension <code>w</code>. Each thread processes one row of <code>B</code> and the entire <code>C</code> vector. <pre><code class="language-c">__global__ void MatrixVectorMulKernel(float* A, float* B, float* C, int w, int h) { int row = blockIdx.x*blockDim.x + threadIdx.x; if (row < h) { float Pvalue = 0; for (int col = 0; col < w; ++col) { Pvalue += B[row*w + col] * C[col]; } A[row] = Pvalue; } } void matrixVectorMul(float* A, float* B, float* C, int w, int h) { int sizeA = h * sizeof(float); int sizeB = h * w * sizeof(float); int sizeC = w * sizeof(float); float *A_d, *B_d, *C_d; cudaMalloc((void **) &A_d, sizeA); cudaMalloc((void **) &B_d, sizeB); cudaMalloc((void **) &C_d, sizeC); cudaMemcpy(B_d, B, sizeB, cudaMemcpyHostToDevice); cudaMemcpy(C_d, C, sizeC, cudaMemcpyHostToDevice); int block_dim = 256; int grid_dim = (h + block_dim - 1) / block_dim MatrixVectorMulKernel<<<grid_dim, block_dim>>>(A_d, B_d, C_d, w, h); cudaMemcpy(A, A_d, sizeA, cudaMemcpyDeviceToHost); cudaFree(A_d); cudaFree(B_d); cudaFree(C_d); } </code></pre> <blockquote> <ol start="3"> <li>Consider the following CUDA kernel and the corresponding host function that calls it:</li> </ol> </blockquote> <pre><code class="language-c">__global__ void foo_kernel(float* a, float* b, unsigned int M, unsigned int N) { unsigned int row = blockIdx.y * blockDim.y + threadIdx.y; unsigned int col = blockIdx.x * blockDim.x + threadIdx.x; if (row < M && col < N) { b[row * N + col] = a[row * N + col] / 2.1f + 4.8f; } } void foo(float* a_d, float* b_d) { unsigned int M = 150; unsigned int N = 300; dim3 bd(16, 32); dim3 gd((N - 1) / 16 + 1, (M - 1) / 32 + 1); foo_kernel<<<gd, bd>>>(a_d, b_d, M, N); } </code></pre> <blockquote> a. What is the number of threads per block? </blockquote> <code>512</code> <blockquote> b. What is the number of threads in the grid? </blockquote> We have <code>gd(19, 5)</code>, so <code>95</code> blocks per grid. <code>95 x 512 = 48,640</code> threads per grid. <blockquote> c. What is the number of blocks in the grid? </blockquote> <code>95</code> <blockquote> d. What is the number of threads that execute the code on line 05? </blockquote> <code>150 x 300 = 45,000</code> <blockquote> <ol start="4"> <li>Consider a 2D matrix with a width of 400 and a height of 500. The matrix is stored as a one-dimensional array. Specify the array index of the matrix element at row 20 and column 10:</li> </ol> </blockquote> <blockquote> a. If the matrix is stored in row-major order. </blockquote> Each row is appended subsequent to each other, so we have <code>20 * 400 + 10 = 8010</code>. <blockquote> b. If the matrix is stored in column-major order. </blockquote> Each column is appended subsequent to each other, so we have <code>10 * 500 + 20 = 5020</code>. <blockquote> <ol start="5"> <li>Consider a 3D tensor with a width of 400, a height of 500, and a depth of 300. The tensor is stored as a one-dimensional array in row-major order. Specify the array index of the tensor element at x=10, y=20, and z=5.</li> </ol> </blockquote> In row-major format, for an array <code>A[z, y, x]</code>, the rightmost index changes fastest. Suppose we have <code>10</code> in each dimension. Then we go: <ul> <li><code>A[0, 0, 0]</code> to <code>A[0, 0, 1]</code> etc.</li> <li>Then <code>A[0, 0, 9]</code> to <code>A[0, 1, 0]</code> when the <code>y</code> dimension changes</li> <li>Then <code>A[0, 9, 9]</code> to <code>A[1, 0, 0]</code> when the <code>z</code> dimension changes</li> </ul> So to locate <code>A[5, 20, 10]</code>: <ul> <li>First we find the start of the relevant plane <code>5 * 500 * 400</code></li> <li>Then within this plane, we find the start of the relevant row <code>20 * 400</code></li> <li>Then within this row, we find the element <code>10</code></li> </ul> So the answer is the sum of all these: <code>1,008,010</code>. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="compute-architecture"><a class="header" href="#compute-architecture">Compute Architecture</a></h1> This chapter elaborates on aspects of the GPU compute architecture that are important for CUDA programmers to reason about. <h2 id="architecture"><a class="header" href="#architecture">Architecture</a></h2> A GPU is organized into an array of Streaming Multiprocessors (SMs). <ul> <li>e.g. A100 has <code>108</code> SMs</li> </ul> Each SM has several processing units called streaming processors or just CUDA cores. These cores share control logic and memory resources <ul> <li>e.g. A100 has <code>64</code> cores per SM, so <code>6,912</code> cores in the GPU</li> </ul> In terms of memory: <ul> <li>Each SM has a set of on-chip memory structures collectively called Memory</li> <li>All SMs also have access to gigabytes of off-chip device memory, called Global Memory</li> <li>For recent GPUs, these are HBM or HBM2</li> <li>We will refer to this as DRAM (dynamic random access memory)</li> </ul> <h2 id="block-scheduling"><a class="header" href="#block-scheduling">Block Scheduling</a></h2> When a kernel is called, a grid of threads is launched. These threads are assigned to SMs on a block by block basis, i.e. all threads in a block are simultaneously assigned to the same SM. Usually, multiple blocks are simultaneously assigned to one SM. The number of blocks that can be processed at the same time by an SM depends on hardware constraints. This means that there is a limit on the number of blocks that are simultaneously executing on a CUDA GPU. The runtime system needs to maintain a list of blocks to execute and assign new blocks to SMs that have finished execution. Since we are guaranteed that all threads on the same block are assigned to the same SM, it is possible for these threads on the same block to interact with each other. One interaction is barrier synchronization. <h2 id="synchronization"><a class="header" href="#synchronization">Synchronization</a></h2> There is a <code>__syncthreads()</code> function to coorindate activities between threads in the same block. When a thread calls it, it will be held at the program location of the call until every thread in the block reaches that location. Each <code>__syncthreads</code> call must be executed by all threads in a block. If we use control flows like <code>if-else</code>, each <code>__syncthreads</code> call is considered a different barrier. Hence, if some threads go through the <code>if</code> branch and others the <code>else</code> branch, we can end in a deadlock as we never reach the state where all threads hit the same <code>__syncthreads</code> call. Note that by design, threads in different blocks cannot perform barrier synchronization with each other. This means that CUDA runtime can execute blocks in any order. This flexibility allows the runtime to efficiently allocate blocks to SMs according the GPU capacity. <h2 id="warps-and-simd"><a class="header" href="#warps-and-simd">Warps and SIMD</a></h2> Here we go into further details on thread scheduling within a block. When a block is assigned to an SM, it gets further divided into 32-thread units called warps. Each warp comprises <code>32</code> threads of consecutive <code>threadIdx</code> values. <ul> <li>threads <code>0-31</code> form the first warp</li> <li>threads <code>32-63</code> form the next warp etc.</li> <li>If the number of threads is not a multiple of <code>32</code>, the warp will be padded with extra dormant threads to complete the warp</li> </ul> Blocks are partitioned into warps based on thread indices. If the blockDim is one-dimensional, the partitioning and ordering is trivial. If there are two or three dimensions, the thread indices will be linearized using a row-major layout before partitioning and ordering into warps. An SM is designed to execute all threads in a warp following the Single Instruction, Multiple Data (SIMD) model. This means that at any point in time, a single instruction is fetched and executed for all threads in a warp (for the most part, common exception is conditional branches). <ul> <li>For e.g., the A100 SM has <code>64</code> cores, split into <code>4</code> processing blocks with <code>16</code> cores each. Threads in the same warp get assigned to the same processing block, which fetches one instruction and executes it for all <code>32</code> threads simultaneously.</li> </ul> The advantage of SIMD is that the cost of the control hardware (the instruction fetch/dispatch unit) is shared across many execution units. This allows for a smaller percentage of the hardware to be dedicated to control flow. <h2 id="control-divergence"><a class="header" href="#control-divergence">Control Divergence</a></h2> In the case of control divergence, i.e. <code>if-else</code> statements exist in the kernel code which depends on <code>threadIdx</code>, then the SIMD model breaks down. To accommodate such cases, the runtime makes multiple passes through each branch. For example if we do <code>if (threadIdx.x < 24)</code>, then one pass is done to compute those threads below the threshold, and another pass for those above the threshold. The threads which are not executing at any point in time are just put on hold / held inactive. The threads then reconverge after the conditional statement and continue executing. Another example is to control the number of loops based on data: <pre><code class="language-c">N = a[threadIdx.x]; for (i=0; i < N; ++i) { // do something } </code></pre> The cost we pay for this flexibility is the extra number of passes we need to make. The most common example of control divergence is handling boundary conditions due to the number of data points not being a multiple of the number of threads. For this case, only the last boundary case has control divergence. Hence when the size of data is large, the effects of control divergence is minimal. An important implication of control divergence is that one cannot assume that all threads in a warp have the same execution timing. Hence we must use <code>__syncwarp</code> if we are depending on such behaviour. <h2 id="warp-scheduling-and-latency-tolerance"><a class="header" href="#warp-scheduling-and-latency-tolerance">Warp Scheduling and Latency Tolerance</a></h2> When assigning threads to SMs, we usually assign far more threads than cores available on the SM. A natural question to ask is whether we should just assign the number of threads to be equal to the number of cores at each request? It turns out that having many threads (or warps) waiting around for execution on SMs is an important feature for GPUs to tolerate long latency operations like reading from global memory. When a warp of threads arrives at an SM and it is not ready for execution (as it is waiting for the result of a previously initiated memory access operation), the warp is not selected. Instead the SM will select another warp that is ready for execution. This mechanism allows the SM to "hide" the latency of long latency operations and fill it with meaningful work. This dynamic scheduling of warps is made even more efficient due to zero-overhead scheduling design of GPUs. In CPUs, context switching is more expensive because when switching context from a thread that is waiting to a thread that is ready to make progress, some work needs to be done to save and restore register contents, adding significant overhead. In GPUs, because SMs hold all the execution states of assigned warps in the hardware registers, there is very little overhead in switching between warps. Since GPUs can effectively "hide" the latency of long latency operations, GPUs dedicate less chip area to latency-reduction mechanisms like cache memories and branch prediction mechanisms, allowing them to focus more on floating point execution and large bandwidth memory access. <h2 id="occupancy"><a class="header" href="#occupancy">Occupancy</a></h2> As we saw, over-subscription of warps to SMs is key for tolerating long latency operations. Hence, the goal is to maximize the number of warps assigned to an SM. <blockquote> Occupancy. The ratio of the number of warps assigned to an SM over the maximum number of warps an SM supports is refrred to as occupancy. </blockquote> Due to different limits on resource allocation, an inefficient configuration can lead to under-utilization and sub-optimal occupancy. For example, the A100 has the following limits: <ul> <li>Up to <code>32</code> blocks per SM</li> <li>Up to <code>64</code> warps per SM</li> <li>i.e. up to <code>2,048</code> thread slots per SM</li> <li>Up to <code>1,024</code> threads per block</li> </ul> If we set <code>block_size=512</code>, <code>num_blocks=4</code>, then we get to max out our thread slots at <code>2,048</code> But if we set <code>block_size=32</code>, and maximize the number of blocks at <code>num_blocks=32</code>, we only utilize <code>1,024</code> thread slots. Because we utilized too few threads per block, we end up with only <code>50%</code> occupancy. Anoter scenario for under-utilization is when the number of threads per block is not divisible by the block size. For example, if we choose <code>block_size=768</code>, then the maximum <code>num_blocks=2</code>. This results in <code>1,536 / 2,048 = 75%</code> occupancy. Another factor that may lower occupancy is the impact of register usage. The A100 has a maximum of <code>65,536</code> registers per SM, so to run at full occupancy, we can only support <code>65,536 / 2,048 = 32</code> registers per thread. If our kernel code uses <code>64</code> registers, then the maximum occupancy we can achieve is <code>1,024 / 2,048 = 50%</code> regardless of the block size. <h2 id="querying-device-properties"><a class="header" href="#querying-device-properties">Querying device properties</a></h2> Get number of CUDA devices in the system: <pre><code class="language-c">int devCount; cudaGetDeviceCount(&devCount); </code></pre> Get device properties: <pre><code class="language-c">cudaDeviceProp devProp; for (unsigned int i = 0; i < devCount; i++) { cudaGetDeviceProperties(&devProp, i); // Do something with devProp } </code></pre> Some attributes of <code>devProp</code>: <ul> <li><code>devProp.maxThreadsPerBlock</code> gives the maximum number of threads allowed in a block in the device.</li> <li><code>devProp.multiProcessorCount</code> gives the number of SMs in the device</li> <li><code>devProp.clockRate</code> gives the clock frequency of the device</li> <li><code>devProp.maxThreadsDim[0]</code> gives the maximum number of threads in the x dimension. Use <code>1</code> for y dimension and <code>2</code> for z dimension.</li> <li><code>devProp.maxGridSize[0]</code> similarly gives the maximum number of blocks in the x dimension in the grid.</li> <li><code>devProp.regsPerBlock</code> gives the number of registers that are available in each SM. The name is a minomer - it is usually equivalent to the number of registers per SM as well.</li> <li><code>devProp.warpSize</code> gives the size of a warp.</li> </ul> <h2 id="exercises-2"><a class="header" href="#exercises-2">Exercises</a></h2> <blockquote> <ol> <li>Consider the following CUDA kernel and the corresponding host function that calls it:</li> </ol> </blockquote> <pre><code class="language-c">__global__ void foo_kernel(int* a, int* b) { unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; if (threadIdx.x < 40 || threadIdx.x >= 104) { b[i] = a[i] + 1; } if (i % 2 == 0) { a[i] = b[i] * 2; } for (unsigned int j = 0; j < 5 - (i % 3); ++j) { b[i] += j; } } void foo(int* a_d, int* b_d) { unsigned int N = 1024; foo_kernel<<< (N + 128 - 1) / 128, 128 >>>(a_d, b_d); } </code></pre> <blockquote> a. What is the number of warps per block? </blockquote> <code>128 / 32 = 4</code> <blockquote> b. What is the number of warps in the grid? </blockquote> Number of blocks is <code>1,024 / 128 = 8</code>. So number of warps is <code>8 * 4 = 32</code>. <blockquote> c. For the statement on line 04: <ul> <li>How many warps in the grid are active?</li> <li>How many warps in the grid are divergent?</li> <li>What is the SIMD efficiency (in %) of warp 0 of block 0?</li> <li>What is the SIMD efficiency (in %) of warp 1 of block 0?</li> <li>What is the SIMD efficiency (in %) of warp 3 of block 0?</li> </ul> </blockquote> <ul> <li>Each block has <code>4</code> warps. Amongst the warps, only warp 2 is inactive. So total active warps = <code>3 * 8 = 24</code>.</li> <li>Warps <code>1</code> and <code>3</code> are divergent.</li> <li>Efficiency of warp 0 of block 0 is <code>100%</code></li> <li>Not clear how SIMD efficiency for a divergent warp is defined. Efficiency of warp 1 of block 0 is <code>8 / 32 = 25%</code> on the <code>True</code> pass and <code>75%</code> on the <code>False</code> pass.</li> <li>Efficiency of warp 3 of block 0 is <code>24 / 32 = 75%</code> on the <code>True</code> pass and <code>25%</code> on the <code>False</code> psas.</li> </ul> <blockquote> d. For the statement on line 07: <ul> <li>How many warps in the grid are active?</li> <li>How many warps in the grid are divergent?</li> <li>What is the SIMD efficiency (in %) of warp 0 of block 0?</li> </ul> </blockquote> <ul> <li>All warps are active</li> <li>All warps are divergent</li> <li><code>50%</code></li> </ul> <blockquote> e. For the loop on line 09: <ul> <li>How many iterations have no divergence?</li> <li>How many iterations have divergence?</li> </ul> </blockquote> The number of <code>j</code> loops depends on the result of <code>j < 5 - (i % 3)</code>: <ul> <li><code>j</code> runs from <code>0-4</code></li> <li><code>j</code> runs from <code>0-3</code></li> <li><code>j</code> runs from <code>0-2</code></li> </ul> Hence there is always divergence in any warp. <blockquote> <ol start="2"> <li>For a vector addition, assume that the vector length is 2000, each thread calculates one output element, and the thread block size is 512 threads. How many threads will be in the grid?</li> </ol> </blockquote> <code>2,048</code> <blockquote> <ol start="3"> <li>For the previous question, how many warps do you expect to have divergence due to the boundary check on vector length?</li> </ol> </blockquote> <code>1</code> warp is divergent, <code>1</code> warp is inactive. <blockquote> <ol start="4"> <li>Consider a hypothetical block with 8 threads executing a section of code before reaching a barrier. The threads require the following amount of time (in microseconds) to execute the sections: 2.0, 2.3, 3.0, 2.8, 2.4, 1.9, 2.6, and 2.9; they spend the rest of their time waiting for the barrier. What percentage of the threads’ total execution time is spent waiting for the barrier?</li> </ol> </blockquote> The slowest thread is <code>3.0</code>, so total time is <code>3 * 8 = 24</code>. The sum of execution times is <code>19.9</code>, so total time waiting is <code>17.1%</code>. <blockquote> <ol start="5"> <li>A CUDA programmer says that if they launch a kernel with only 32 threads in each block, they can leave out the __syncthreads() instruction wherever barrier synchronization is needed. Do you think this is a good idea? Explain.</li> </ol> </blockquote> No, because we still need it for syncing across different warps in the same block. Also if there's divergence, there is no guarantee that threads in the same warp execute in lock step. <blockquote> <ol start="6"> <li>If a CUDA device’s SM can take up to 1536 threads and up to 4 thread blocks, which of the following block configurations would result in the most number of threads in the SM?</li> </ol> <ul> <li>128 threads per block</li> <li>256 threads per block</li> <li>512 threads per block</li> <li>1024 threads per block</li> </ul> </blockquote> <code>512</code> with <code>3</code> thread blocks. <blockquote> <ol start="7"> <li>Assume a device that allows up to 64 blocks per SM and 2048 threads per SM. Indicate which of the following assignments per SM are possible. In the cases in which it is possible, indicate the occupancy level.</li> </ol> </blockquote> <ul> <li>8 blocks with 128 threads each <ul> <li>Possible, occupancy = <code>50%</code></li> </ul> </li> <li>16 blocks with 64 threads each <ul> <li>Possible, occupancy = <code>50%</code></li> </ul> </li> <li>32 blocks with 32 threads each <ul> <li>Possible, occupancy = <code>50%</code></li> </ul> </li> <li>64 blocks with 32 threads each <ul> <li>Possible, occupancy = <code>100%</code></li> </ul> </li> <li>32 blocks with 64 threads each <ul> <li>Possible, occupancy = <code>100%</code></li> </ul> </li> </ul> <blockquote> <ol start="8"> <li>Consider a GPU with the following hardware limits: 2048 threads per SM, 32 blocks per SM, and 64K (65,536) registers per SM. For each of the following kernel characteristics, specify whether the kernel can achieve full occupancy. If not, specify the limiting factor.</li> </ol> </blockquote> <ul> <li>The kernel uses 128 threads per block and 30 registers per thread. <ul> <li>Yes, we can hit full occupancy</li> </ul> </li> <li>The kernel uses 32 threads per block and 29 registers per thread. <ul> <li>No, not enough threads per block, we can only get <code>50%</code> occupancy</li> </ul> </li> <li>The kernel uses 256 threads per block and 34 registers per thread. <ul> <li>No, we have too many registers per thread, so we cannot have <code>2,048</code> threads.</li> </ul> </li> </ul> <blockquote> <ol start="9"> <li>A student mentions that they were able to multiply two 1024×1024 matrices using a matrix multiplication kernel with 32×32 thread blocks. The student is using a CUDA device that allows up to 512 threads per block and up to 8 blocks per SM. The student further mentions that each thread in a thread block calculates one element of the result matrix. What would be your reaction and why?</li> </ol> </blockquote> It does not seem possible as he has <code>1,024</code> threads per block which exceeds the hardware limit. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="memory-architecture"><a class="header" href="#memory-architecture">Memory Architecture</a></h1> This chapter focuses on the memory architecture of GPUs and how to organize and position data for efficient access. Consider the matrix multiplication code back in chapter 3. This performs a for loop to compute the dot product between a row of M and a column of N. <pre><code class="language-c">for (int k = 0; k < Width; ++k) { Pvalue += M[row*Width + k] * N[k*Width + col]; } </code></pre> In every iteration of the loop: <ul> <li>We perform two global memory accesses (for M and N respectively)</li> <li>We perform one floating point multiplication and one floating point addition</li> <li>The ratio of floating point operation (FLOP) to bytes accessed is <code>2 FLOP / 8 bytes = 0.25 FLOP/B</code>. This is called arithmetic intensity or computational intensity in the literature.</li> </ul> The compute to global memory access has implications for the performance of a GPU kernel. <ul> <li>An A100 has peak global memory bandwidth of <code>1555 GB/s</code>. With <code>0.25 FLOP/B</code>, we can get <code>1555 GB/s * 0.25 FLOP/B = 389 giga FLOPs/s (GLOPS)</code> number of single precision floating point operations.</li> <li>But A100 has a cap of <code>19,500 GLOPS</code>, so we are only using <code>2%</code> of the compute capacity!</li> <li>We are severely limited by memory accesses, and this type of program is called <code>memory-bound</code> programs</li> </ul> <h2 id="roofline-model"><a class="header" href="#roofline-model">Roofline Model</a></h2> The roofline model provides a visual representation of the performance of an application relative to hardware limits. <div class="table-wrapper"><table><thead><tr><th style="text-align: center"><img src="pmpp/../images/pmpp/ch_05_roofline_model.png" alt="Roofline model illustrating bandwidth and compute ceilings" /></th></tr></thead><tbody> <tr><td style="text-align: center">Roofline model for analyzing GPU efficiency</td></tr> </tbody></table> </div> The axes meanings: <ul> <li><code>x-axis</code> is the computational intensity measured in <code>FLOP/B</code>. This quantity is determined by how the application is written.</li> <li><code>y-axis</code> is the computational throughput, i.e. obtained by multiplying the computational intensity by the maximum memory bandwith of our hardware device (a fixed constant) <ul> <li>The peak throughput horizontal line is the hard cap of the hardware on the peak throughput it can support</li> </ul> </li> </ul> The sloping line labelled peak bandwidth represents the hardward limit of the memory bandwidth. The slope is obtained by multipling the current position on the x-axis by the maximum memory bandwidth (say <code>1555GB/s</code> in the example above), which gives us a slope. We cannot have any point above this line because it would require a memory bandwidth that exceeds the hardware limits. A point on the plot represents an application with its computational intensity (FLOPS / byte accessed) and computational throughput (FLOPS / s). Note that in real applications, these are measured using programs like NVIDIA NSight using hardware counters. The position of the point relative to the lines tells us its efficiency: <ul> <li>Points can only be below the two lines, as we cannot exceed hardware limits</li> <li>Points closer to the lines are more efficient, as we are nearing hardware limits</li> </ul> Based on the diagram: <ul> <li>Point A3 is compute bound - it is already quite computationally intensive and near the peak computational throughput, so we need a more powerful GPU to improve it further.</li> <li>Point A2 is inefficient - neither memory nor compute bound. It can be improved by improving the memory bandwidth utilization.</li> <li>Point A1 is memory bound - it is utilizing memory bandwidth effectively but the compute intensity can be improved. e.g. To fully utilize the <code>19,500 GLOPS</code> from the A100, we need to increase compute efficiency to <code>19,500 (GFLOP/s) / 1,555 (GB/s) = 12.5 FLOP/B</code>. This means that for each single precision floating point number accessed, we need to perform <code>50 FLOPS</code>! This can be achieved for certain applications like matrix multiplication.</li> </ul> <h2 id="cuda-memory-types"><a class="header" href="#cuda-memory-types">CUDA Memory Types</a></h2> There are different memory types on a GPU with different uses: <ul> <li>Registers <ul> <li>Registers are on-chip memory, which means very short latency, high bandwidth access</li> <li>Highly limited number of registers per SMs, using too many in kernel code can result in low occupancy</li> <li>Allocated to individual threads - each thread can only access its own registers</li> </ul> </li> <li>Shared memory <ul> <li>Also on-chip memory, but it resides outside of the Processing Unit</li> <li>Accessing data requires a memory load (as opposed to registers), so it is slower than registers, but the memory load is much faster than global memory as it is on-chip</li> <li>Allocated to thread blocks - all threads in a block can access shared memory variables</li> </ul> </li> <li>Global memory <ul> <li>All threads can access</li> <li>Off-chip memory, so relatively high latency and low bandwidth access</li> </ul> </li> <li>Local memory <ul> <li>Placed in global memory and has similar access latency</li> <li>Not shared across threads</li> <li>Each thread has its own section of global memory that it uses as its own private local memory that is private to the thread but cannot be allocated in registers</li> <li>e.g. statically allocated arrays</li> </ul> </li> </ul> Efficient GPU code requires us to make full use of registers whenever possible: <ul> <li>Access bandwidth of registers is at least 2 orders of magnitude higher than global memory</li> <li>When a variable is stored in registers, its access no longer consume off-chip global memory bandwith, which is reflected as an increased compute to global memory access ratio</li> <li>Each access to registers involves fewer instructions than access to global memory due to built-in register operands. <ul> <li>e.g. <code>fadd r1, r2, r3</code> directly adds values in registers <code>r2, r3</code> and stores it into <code>r1</code>. There is no loading step</li> <li>But if the first operand is in the global memory, we need to do an extra loading step</li> </ul> <pre><code>load r2, r4, offset fadd r1, r2, r3 </code></pre> </li> <li>The energy consumed for accessing a value in the register file is an order of magnitude lower than accessing the value from global memory.</li> </ul> <h2 id="cuda-syntax"><a class="header" href="#cuda-syntax">CUDA Syntax</a></h2> We can declare variables and reside them in the intended type of memory based on our application design. <blockquote> Scope. Scope identifies the set of threads than can access a variable. If the scope is a single thread, a private version of the variable is created for each thread. Lifetime. The lifetime is the portion of a program's execution duration when the variable is available, If the lifetime is within a grid's execution, it must be declared within the kernel function. When the kernel is executed multiple times, the value is not maintained across these invocations. If the lifetime is entire application, it must be declared outside of any function body. </blockquote> Automatic Scalar Variables. All automatic scalar variables declared in the kernel and device functions are automatically placed into registers. This means that a private copy of that variable in generated for every thread that executes the kernel function. <ul> <li>Scope: Thread</li> <li>Lifetime: Grid</li> </ul> Automatic array variables. These are array variables defined in kernel code, such as <code>float buf[128]</code>. Note that the array size must be specified at compile time. These are put into local memory (part of global memory), so access is very slow. In practice, these are seldom used. <ul> <li>Scope: Thread</li> <li>Lifetime: Grid</li> </ul> Shared Memory variables. These are declared in the kernel function and put into shared memory. The scope of a shared variable is in a thread block, so all threads in a block see the same version of a shared variable (and can read and write to it). A private version of the variable is created for each block during kernel execution. Accessing shared variables from the shared memory is extremely fast and highly parallel, and it is a good way for threads in a block to communicate with each other. <ul> <li>Declaration: <code>__device__ __shared__ int SharedVar;</code></li> <li>Scope: Block</li> <li>Lifetime: Grid</li> </ul> Constant Memory variables. These are declared in the host code outside of any function body, and its values cannot be changed by kernel function code. The name is indicative: its usually used for global constants like model parameters or for things like a fixed convolution filter. Constant variables are stored in global memory but are cached for extremely fast and parallel access. The total size of constant variables in an application is limited to <code>65,536 bytes</code>. <ul> <li>Declaration: <code>__device__ __constant__ int ConstVar;</code></li> <li>Scope: Grid</li> <li>Lifetime: Application</li> </ul> Global Memory variables. These are placed in global memory where access is slow. The advantage of global memory is that it is accessible by all threads and the contents persist throughout the entire application. So it is a way for threads to collaborate across blocks. However, because there is no barrier synchronization method for threads across blocks, usually global variables are used to pass information from one kernel invocation to another kernel. <ul> <li>Declaration: <code>__device__ int GlobalVar;</code></li> <li>Scope: Grid</li> <li>Lifetime: Application</li> </ul> <h2 id="tiling"><a class="header" href="#tiling">Tiling</a></h2> The intrinsic memory tradeoff in CUDA is that global memory is large but slow, and shared memory is small but fast. Hence a good strategy is to partition the data into subsets called tiles so that each tile fits into the shared memory. Consider the matrix multiplication example in chapter 3 where each thread computes one value of the output matrix <code>P = M dot N</code>, and thus access one row of <code>M</code> and one column of <code>N</code>. <ul> <li>For brevity, define My,x as <code>M[y*Width + x]</code>, Ny,x as <code>N[y*Width + x]</code>, Py,x as <code>P[y*Width + x]</code></li> <li>Suppose <code>Width=4</code> and our block size is <code>2 x 2</code></li> </ul> Now observe that there are duplicate global memory accesses across threads in a block: <ul> <li>M0,0 is accessed for computing both P0,0 and P0,1</li> <li>N0,0 is accessed for computing both P0,0 and P1,0</li> <li>And so on</li> </ul> Instead of accessing global memory twice, we could have accessed it once, loaded it into shared memory, and then have threads read from shared memory thereafter (almost negligible cost compared to accessing global memory). In fact, a given value (like M0,0) will be accessed <code>block_size</code> times for a given block, so the larger our block size, the more efficient our memory access can be. We are limited by the size of shared memory, so we need to further tile <code>M</code> and <code>N</code> matrices so that it will fit into shared memory. For simplicity, consider tiling <code>M</code> and <code>N</code> into <code>2 x 2</code> tiles also. Our strategy will now be like so: <ul> <li>Initialize <code>2 x 2</code> shared variables <code>Mds</code> and <code>Nds</code></li> <li>Load a tile of <code>M</code> into <code>Mds</code>: <ul> <li>For e.g. the first tile would be M0,0,M0,1,M1,0,M1,1</li> </ul> </li> <li>Load a tile of <code>N</code> into <code>Nds</code>: <ul> <li>For e.g. the first tile would be N0,0,N0,1,N1,0,N1,1</li> </ul> </li> <li>Launch threads to accumulate partial dot product into thread-specific <code>Pvalue</code>: <ul> <li>Instead of reading from <code>M</code> and <code>N</code>, we read the required variables from <code>Mds</code> and <code>Nds</code> instead</li> <li>e.g. thread0,0 computes Pvalue0,0 += Mds0,0∗Nds0,0+Mds0,1∗Nds1,0</li> <li>Note that this is a partial dot product for Pvalue0,0, we need to accumulate to it again</li> </ul> </li> <li>Move on to the next tile and accumulate to <code>Pvalue</code> again: <ul> <li>e.g. Second tile is M0,2,M0,3,M1,2,M1,3</li> <li>e.g. Second tile is N2,0,N2,1,M3,0,M3,1</li> </ul> </li> </ul> Note that with this tiling strategy, we will need to do this in multiple steps (or phases) for a given block. In general, we need <code>Width / TILE_WIDTH</code> phases to complete the matrix multiplication. <h2 id="a-tiled-matmul-kernel"><a class="header" href="#a-tiled-matmul-kernel">A tiled matmul kernel</a></h2> Now we can tackle the matmul kernel. <pre><code class="language-c"># define TILE_WIDTH 16 __global__ void matrixMulKernel(float* M, float* N, float* P, int Width) { __shared__ float Mds[TILE_WIDTH][TILE_WIDTH]; __shared__ float Nds[TILE_WIDTH][TILE_WIDTH]; int bx = blockIdx.x; int by = blockIdx.y; int tx = threadIdx.x; int ty = threadIdx.y; // Identify P[Row, Col] to work on int Row = by * TILE_WIDTH + ty; int Col = bx * TILE_WIDTH + tx; float Pvalue = 0; for (int ph = 0; ph < Width/TILE_WIDTH; ++ph) { // Load M, N into shared memory Mds[ty][tx] = M[Row*Width + ph*TILE_WIDTH + tx]; Nds[ty][tx] = N[(ph*TILE_WIDTH + ty)*Width + Col]; __syncthreads(); for (int k = 0; k < TILE_WIDTH; ++k) { Pvalue += Mds[ty][k] * Nds[k][tx]; } __syncthreads(); } P[Row*Width + Col] = Pvalue; } </code></pre> Notes on the above: <ul> <li><code>Mds, Nds</code> are shared variables that are common for all threads in a block</li> <li>Since each thread is working on one element of output <code>P</code>, we find <code>Row</code> and <code>Col</code> respectively (to compute <code>P[Row, Col]</code>) <ul> <li>Note that the algorithm assumes each block computes exactly one tile, so <code>TILE_WIDTH = BLOCK_SIZE</code></li> <li>Therefore <code>Row = by * TILE_WIDTH + ty</code></li> <li>And <code>Col = bx * TILE_WIDTH + tx</code></li> </ul> </li> </ul> Now we tackle the for loop part which loops through the phase counter <code>ph</code>. <ul> <li>Recall that each phase uses one tile of <code>M</code> and one tile of <code>N</code>.</li> <li>The phase counter <code>ph</code> tells us how many phases we have already processed</li> </ul> <h3 id="loading-data-into-shared-memory"><a class="header" href="#loading-data-into-shared-memory">Loading Data into Shared Memory</a></h3> For each phase, we first load data into shared memory: <ul> <li>Each thread loads one value of <code>M</code> nad one value of <code>N</code> into shared memory</li> <li>Since we have <code>TILE_WIDTH x TILE_WIDTH</code> number of threads per block, in each phase we load this many elements into <code>Mds</code> and <code>Nds</code> respectively</li> </ul> Which data should each thread load? <ul> <li>For <code>M</code>: <ul> <li>Note that <code>Row</code> already gives us the row of <code>M</code> to be loaded for this thread (constant across phases)</li> <li>The beginning column for this block and phase is <code>ph * TILE_WIDTH</code>, because we have processed <code>ph</code> number of tiles to the left</li> <li>So we want <code>M[Row, ph*TILE_WIDTH]</code>, which is indexed as <code>M[Row*Width + ph*TILE_WIDTH + tx]</code> in row major form</li> </ul> </li> <li>For <code>N</code>: <ul> <li><code>Col</code> already gives us the column of <code>N</code> to be loaded for this thread (constant across phases)</li> <li>The beginning row for this block and phase is <code>ph * TILE_WIDTH</code></li> <li>So we want <code>N[ph*TILE_WIDTH + ty, Col]</code>, which is indexed as <code>N[(ph*TILE_WIDTH + ty)*Width + Col]</code></li> </ul> </li> </ul> Finally, we call <code>__syncthreads()</code> barrier synchronization. This is necessary because for the computation phase, each thread requires values of <code>Mds,Nds</code> that other threads have loaded for it. Hence we must make sure all threads have loaded into shared memory before we proceed to compute. <ul> <li>This is called read-after-write synchronization - we require data to be written before we can read from it</li> <li>Also called true dependence because there is truly a dependence from the reader on the written data</li> </ul> <h3 id="computation-phase"><a class="header" href="#computation-phase">Computation Phase</a></h3> Now we have loaded the requisite data into <code>Mds</code> and <code>Nds</code>, we can accumulate into <code>Pvalue</code>. The tiled matmul is such that we accumulate into <code>Pvalue</code> as though we were just doing a simple dot product between <code>Mds</code> and <code>Nds</code>. This compute phase is very fast as we are reading from shared memory. Finally, another <code>__syncthreads()</code> is called. This is to ensure that computation has completed before we start loading values into <code>Mds</code> and <code>Nds</code> in the next phase, overwriting their values. <ul> <li>This is called write-after-read synchronization - threads must wait for data to be read by all threads that need it before overwriting it</li> <li>Another name is false dependence. The dependence is false because the writing thread actually does not need any data from the reading thread - the dependence occurs purely because they are sharing the same memory location.</li> </ul> <h3 id="summary-2"><a class="header" href="#summary-2">Summary</a></h3> The savings from tiled matmul is significant. With <code>16 x 16</code> tiles, we reduce global memory access by a factor of <code>16</code>, which brings our <code>0.25 FLOP/B</code> to <code>4 FLOP/B</code>. There are further optimizations we can make later. <h2 id="boundary-checks"><a class="header" href="#boundary-checks">Boundary Checks</a></h2> In the previous section, we assume that <code>Width</code> is a multiple of <code>TILE_WIDTH</code>. This section extends it to handle arbitrary matrix sizes. As with before, we need to add boundary checks to avoid accessing non-existent values or even accessing wrong values that corrupt the final output. The modified kernel code is shown (only the for loop part): <pre><code class="language-c">float Pvalue = 0; for (int ph = 0; ph < ceil(Width/(float) TILE_WIDTH); ++ph) { if ((Row < Width) && (ph*TILE_WIDTH+tx) < Width) Mds[ty][tx] = M[Row*Width + ph*TILE_WIDTH + tx]; else Mds[ty][tx] = 0.0f; if ((ph*TILE_WIDTH+ty) < Width && Col < Width) Nds[ty][tx] = N[(ph*TILE_WIDTH+ty)*Width + Col]; else Nds[ty][tx] = 0.0f; __syncthreads(); for (int k = 0; k < TILE_WIDTH; ++k) { Pvalue += Mds[ty][k] * Nds[k][tx]; } __syncthreads(); } if (Row < Width) && (Col < Width) P[Row*Width + Col] = Pvalue; </code></pre> For boundary checks on <code>M</code>: <ul> <li>As we recall, the <code>y</code> and <code>x</code> element accessed are <code>Row</code> and <code>ph*TILE_WIDTH + tx</code> respectively</li> <li>Hence we just need to check if both of these are less than <code>Width</code></li> </ul> For boundary checks on <code>N</code>: <ul> <li>The <code>y</code> and <code>x</code> element accessed are <code>(ph*TILE_WIDTH+ty)</code> and <code>Col</code> respectively</li> <li>Hence we just need to check if both of these are less than <code>Width</code></li> </ul> For accessing <code>M</code> and <code>N</code>, if we are accessing an out of bounds value, simply load <code>0.0</code> into <code>Mds</code> and <code>Nds</code> respectively. This will ensure that when accumulating into <code>Pvalue</code>, using an out of bound value will result in adding <code>0.0</code> and hence not affect the final output. Finally when assigning values into <code>P</code>, we just check if <code>Row, Col</code> are within bounds. <h2 id="impact-of-memory-use-on-occupancy"><a class="header" href="#impact-of-memory-use-on-occupancy">Impact of Memory Use on Occupancy</a></h2> Earlier on we saw that register usage can reduce occupancy. Similarly, using too much shared memory can reduce the number of blocks assigned and hence reduce occupancy. For example, an A100 can have up to <code>164KB</code> of shared memory per SM and support a maximum of <code>2048</code> threads per SM. <ul> <li>This means that to use up all thread slots, a thread block cannot use more than an average of <code>164 KB / 2048 threads = 82B / thread</code></li> </ul> For our tiled matmul example: <ul> <li>Each block has <code>TILE_WIDTH * TILE_WIDTH</code> threads per block</li> <li>Each block uses <code>4B * TILE_WIDTH * TILE_WIDTH</code> shared memory for <code>Mds</code> and <code>Nds</code> respectively</li> <li>So it uses <code>8B</code> per thread, which is less than <code>82B</code> and we are fine</li> </ul> But consider a kernel that uses <code>32KB</code> of shared memory, and each block has <code>256</code> threads. <ul> <li>Then each thread uses on average <code>32KB / 256 = 128B</code>. This exceeds the <code>82B/thread</code> limit.</li> <li>So we can only launch at most <code>164KB / 128B = 1.3k threads</code>, which is around <code>60%</code> occupancy</li> </ul> <h3 id="dynamic-sizing-of-shared-memory"><a class="header" href="#dynamic-sizing-of-shared-memory">Dynamic sizing of shared memory</a></h3> Since the size of shared memory on an SM is device specific, we want the host code to dynamically adjust the <code>TILE_WIDTH</code> and hence amount of shared memory used by a kernel. <ul> <li>This can be done using <code>cudaGetDeviceProperties</code>, and using <code>devProp.sharedMemPerBlock</code> to give the amount of shared memory available in each SM.</li> </ul> Unfortunately, we cannot specify <code>Mds</code> and <code>Nds</code> as before using <code>TILE_WIDTH</code>, because the value of <code>TILE_WIDTH</code> must be available at compile time. If we want to change this value, we must recompile the code. The workaround for such a scenario is to use <code>extern __shared__ Mds_Nds[];</code> inside the kernel code. This allows for a dynamically allocated shared array. <ul> <li>However, since we only have one merged array, we'll need to manually define where <code>Mds</code> ends and <code>Nds</code> starts</li> <li>We also need to linearize our access to the array</li> </ul> At runtime, we can dynamically configure the amount of shared memory for each block using the third configuration parameter for a kernel call: <pre><code class="language-c">size_t size = calculate_appropriate_SM_usage(devProp.sharedMemPerBlock, ...); matrixMulKernel<<<dimGrid,dimBlock,size>>>(Md, Nd, Pd, Width, size/2, size/2); </code></pre> Notes: <ul> <li>Note that <code>size_t</code> is a built-in type for declaring a variable to hold the size information for dynamically allocated data structures, in number of bytes.</li> <li>Also note that we pass <code>size/2</code> for the size of <code>Mds</code> and <code>Nds</code> respectively</li> </ul> The kernel code can then be modified to accommodate the new structure: <pre><code class="language-c">#define TILE_WIDTH 16 __global__ void matrixMulKernel(float* M, float* N, float* P, int Width, unsigned Mdz_sz, unsigned Nds_sz) { extern __shared__ char float Mds_Nds[]; float *Mds = (float*) Mds_Nds; float *Nds = Mds_Nds + Nds_sz; } </code></pre> Note that we would need to use linearized accesses to <code>Mds</code> and <code>Nds</code>. e.g. <code>Mds[ty][tx]</code> becomes <code>Mds[ty*TILE_WIDTH+tx]</code>. <h2 id="exercises-3"><a class="header" href="#exercises-3">Exercises</a></h2> <blockquote> <ol> <li>Consider matrix addition. Can one use shared memory to reduce the global memory bandwidth consumption? Hint: Analyze the elements that are accessed by each thread and see whether there is any commonality between threads.</li> </ol> </blockquote> No - each element of input matrices <code>A + B</code> are accessed exactly once, so no memory savings. <blockquote> <ol start="2"> <li>Draw the equivalent of Fig. 5.7 for a 8×8 matrix multiplication with 2×2 tiling and 4×4 tiling. Verify that the reduction in global memory bandwidth is indeed proportional to the dimension size of the tiles.</li> </ol> </blockquote> <ul> <li>With 2x2 tiling, there are 4 phases, so each element gets accessed from global memory 4x</li> <li>With 4x4 tiling, there are 2 phases, so each element gets accessed from global memory 2x</li> <li>So the reduction in memory bandwidth is proportional to the tile size</li> </ul> With 2x2 tiling, each element gets accessed twice per phase. With 4x4 tiling, each element gets accessed 4x per phase. So the reduction in global memory bandwidth is scaled accordingly. <blockquote> <ol start="3"> <li>What type of incorrect execution behavior can happen if one forgot to use one or both __syncthreads() in the kernel of Fig. 5.9?</li> </ol> </blockquote> This is the tiled matrix multiplication code. <ul> <li>If we miss the first <code>__syncthreads</code>, we could get "read-before-write" situations where another thread tries to load the values of <code>Mds</code> or <code>Nds</code> before it is written. This leads to wrong outputs.</li> <li>If we miss the second <code>__syncthreads</code>, we could get "write-before-read" situations, where another thread has already overwritten the values of <code>Mds</code> or <code>Nds</code> in the next phase before we have read from it. This leads to wrong outputs.</li> </ul> <blockquote> <ol start="4"> <li>Assuming that capacity is not an issue for registers or shared memory, give one important reason why it would be valuable to use shared memory instead of registers to hold values fetched from global memory? Explain your answer.</li> </ol> </blockquote> Threads in the same block can access the shared memory, so we can benefit as long as our application requires multiple threads to access the same values. Registers are private to each thread and we cannot benefit. <blockquote> <ol start="5"> <li>For our tiled matrix-matrix multiplication kernel, if we use a 32×32 tile, what is the reduction of memory bandwidth usage for input matrices M and N?</li> </ol> </blockquote> 32x <blockquote> <ol start="6"> <li>Assume that a CUDA kernel is launched with 1000 thread blocks, each of which has 512 threads. If a variable is declared as a local variable in the kernel, how many versions of the variable will be created through the lifetime of the execution of the kernel?</li> </ol> </blockquote> 512,000 <blockquote> <ol start="7"> <li>In the previous question, if a variable is declared as a shared memory variable, how many versions of the variable will be created through the lifetime of the execution of the kernel?</li> </ol> </blockquote> 1000 <blockquote> <ol start="8"> <li>Consider performing a matrix multiplication of two input matrices with dimensions N×N. How many times is each element in the input matrices requested from global memory when: a. There is no tiling? b. Tiles of size T×T are used?</li> </ol> </blockquote> Easier to reason about total number of accesses: <ul> <li>With no tiling, each thread accesses <code>2N</code> elements and there are <code>NxN</code> threads, so <code>2N^3</code></li> <li>With tiling, each block accesses <code>2TN</code> elements, and there are <code>N/T * N/T</code> blocks, so <code>2N^3/T</code> accesses</li> </ul> <blockquote> <ol start="9"> <li>A kernel performs 36 floating-point operations and seven 32-bit global memory accesses per thread. For each of the following device properties, indicate whether this kernel is compute-bound or memory-bound. a. Peak FLOPS=200 GFLOPS, peak memory bandwidth=100 GB/second b. Peak FLOPS=300 GFLOPS, peak memory bandwidth=250 GB/second</li> </ol> </blockquote> The computational intensity is <code>36 FLOP / (7 * 4B) = 1.3 FLOP/B</code> <ul> <li>With peak memory bandwidth of <code>100 GB/s</code>, our throughput is <code>1.3 (FLOP/B) * (100 GB/s) = 130 GFLOPS</code>. This means we are memory bound as we are not hitting peak FLOPS</li> <li>With peak memory bandwidth of <code>250 GB/s</code>, our throughput is <code>1.3 (FLOP/B) * (250 GB/s) = 325 GLOPS</code>. This means we are compute bound as we have hit peak FLOPS</li> </ul> <blockquote> <ol start="10"> <li>To manipulate tiles, a new CUDA programmer has written a device kernel that will transpose each tile in a matrix. The tiles are of size BLOCK_WIDTH by BLOCK_WIDTH, and each of the dimensions of matrix A is known to be a multiple of BLOCK_WIDTH. The kernel invocation and code are shown below. BLOCK_WIDTH is known at compile time and could be set anywhere from 1 to 20.</li> </ol> </blockquote> <pre><code class="language-cpp">01 dim3 blockDim(BLOCK_WIDTH, BLOCK_WIDTH); 02 dim3 gridDim(A_width/blockDim.x, A_height/blockDim.y); 03 BlockTranspose<<<gridDim, blockDim>>>(A, A_width, A_height); 04 __global__ void 05 BlockTranspose(float* A_elements, int A_width, int A_height) 06 { 07 __shared__ float blockA[BLOCK_WIDTH][BLOCK_WIDTH]; 08 int baseIdx = blockIdx.x * BLOCK_WIDTH + threadIdx.x; 09 baseIdx += (blockIdx.y * BLOCK_WIDTH + threadIdx.y) * A_width; 10 blockA[threadIdx.y][threadIdx.x] = A_elements[baseIdx]; 11 A_elements[baseIdx] = blockA[threadIdx.x][threadIdx.y]; 12 } </code></pre> <blockquote> a. Out of the possible range of values for BLOCK_WIDTH, for what values of BLOCK_WIDTH will this kernel function execute correctly on the device? </blockquote> The code works like so: <ul> <li>Each thread handles an element of <code>A</code></li> <li>The tile / block size is <code>BLOCK_WIDTH</code></li> <li>For each thread: <ul> <li><code>threadIdx.x</code> and <code>threadIdx.y</code> give the <code>x,y</code> coordinates within a block</li> <li>We wish to swap these elements by temporarily assigning into <code>blockA</code></li> </ul> </li> <li>What is the index of the element we are accessing? <ul> <li>The <code>y</code> index in <code>A</code> would be <code>blockIdx.y * BLOCK_WIDTH + threadIdx.y</code></li> <li>The <code>x</code> index in <code>A</code> would be <code>blockIdx.x * BLOCK_WIDTH + threadIdx.x</code></li> <li>So the serialized form is <code>y*A_width + x</code> as shown in the code</li> </ul> </li> </ul> The problem with the code is that we have a potential "write-before-read" situation in line 10-11: <ul> <li>Consider <code>thread1: threadIdx.y=i, threadIdx.x=j</code> and <code>thread2: threadIdx.y=j, threadIdx.x=i</code> which are two threads in mirror coordinates to each other</li> <li>It is possible that <code>thread1</code> completes the transpose operation before <code>thread2</code> and writes into <code>A_elements</code> before <code>thread2</code> reads</li> <li>Then, we end up with <code>thread1</code> and <code>thread2</code> writing the same value into their designated slots which is incorrect</li> </ul> The only <code>BLOCK_WIDTH</code> value which works is <code>1</code>. <blockquote> b. If the code does not execute correctly for all BLOCK_SIZE values, what is the root cause of this incorrect execution behavior? Suggest a fix to the code to make it work for all BLOCK_SIZE values. </blockquote> Add a <code>__syncthreads()</code> between lines 10 and 11. <blockquote> <ol start="11"> <li>Consider the following CUDA kernel and the corresponding host function that calls it:</li> </ol> </blockquote> <pre><code class="language-cpp">01 __global__ void foo_kernel(float* a, float* b) { 02 unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; 03 float x[4]; 04 __shared__ float y_s; 05 __shared__ float b_s[128]; 06 for (unsigned int j = 0; j < 4; ++j) { 07 x[j] = a[j * blockDim.x * gridDim.x + i]; 08 } 09 if (threadIdx.x == 0) { 10 y_s = 7.4f; 11 } 12 b_s[threadIdx.x] = b[i]; 13 __syncthreads(); 14 b[i] = 2.5f * x[0] + 3.7f * x[1] + 6.3f * x[2] + 8.5f * x[3] 15 + y_s * b_s[threadIdx.x] + b_s[(threadIdx.x + 3) % 128]; 16 } 17 18 void foo(int* a_d, int* b_d) { 19 unsigned int N = 1024; 20 foo_kernel<<< (N + 128 - 1) / 128, 128 >>>(a_d, b_d); 21 } </code></pre> <blockquote> a. How many versions of the variable i are there? </blockquote> <code>i</code> is private to each thread, so <code>1024</code> copies. <blockquote> b. How many versions of the array x[] are there? </blockquote> <code>x[4]</code> is a local memory variable and private to each thread, so <code>1024</code> copies. <blockquote> c. How many versions of the variable y_s are there? </blockquote> <code>y_s</code> is a shared memory variable, one per block, so <code>8</code> copies. <blockquote> d. How many versions of the array b_s[] are there? </blockquote> <code>b_s</code> is also shared memory, so <code>8</code> copies. <blockquote> e. What is the amount of shared memory used per block (in bytes)? </blockquote> We have <code>(128 + 1) * 4B = 516B</code>. <blockquote> f. What is the floating-point to global memory access ratio of the kernel (in OP/B)? </blockquote> There are <code>6</code> memory accesses, so <code>24B</code> in each thread: <ul> <li>Line 7 has <code>4</code> loads</li> <li>Line 12 has <code>1</code> load</li> <li>Line 14 has <code>1</code> store</li> </ul> There are <code>10</code> floating point operations on lines 14 and 15: <ul> <li><code>5</code> multiplications</li> <li><code>5</code> additions</li> </ul> So FLOP/B is <code>10 / 24 = 0.42 FLOP/B</code>. <blockquote> <ol start="12"> <li>Consider a GPU with the following hardware limits: 2048 threads/SM, 32 blocks/SM, 64K (65,536) registers/SM, and 96 KB of shared memory/SM. For each of the following kernel characteristics, specify whether the kernel can achieve full occupancy. If not, specify the limiting factor.</li> </ol> a. The kernel uses 64 threads/block, 27 registers/thread, and 4 KB of shared memory/block. </blockquote> Checks: <ul> <li>Total threads: <code>64 threads/block * 32 blocks = 2048</code>. Hits thread limits ✔</li> <li>Shared memory: given shared memory limits, we can only have <code>96KB / 4KB = 24 blocks</code>. So occupancy can only be 75%. ✖</li> <li>Registers: <code>27 * 2048 = 55,296</code>. Within limits ✔</li> </ul> Cannot reach full occupancy due to shared memory usage. <blockquote> b. The kernel uses 256 threads/block, 31 registers/thread, and 8 KB of shared memory/block. </blockquote> Checks: <ul> <li>Total threads: <code>256 threads/block * 8 blocks = 2048</code>. Hits thread limits ✔</li> <li>Shared memory: given shared memory limits, we have <code>96KB / 8KB = 12 blocks</code>. So we can have 8 blocks no problem. ✔</li> <li>Registers: <code>31 * 2048 = 63,488</code>. Within limits ✔</li> </ul> Can reach full occupancy. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="performance-considerations"><a class="header" href="#performance-considerations">Performance Considerations</a></h1> This chapter goes deeper into off chip memory (DRAM) and discusses memory coalescing and memory latency hiding. It also discusses thread granularity coarsening. <h2 id="memory-coalescing"><a class="header" href="#memory-coalescing">Memory Coalescing</a></h2> One of the most important factors in CUDA kernel performance is accessing data in the global memory. This chapter further discusses memory coalescing techniques to move data between global memory and shared memory or registers efficiently. Global memory is implemented with DRAM: <ul> <li>Access latency to a bit in DRAM is relatively slow (tens of nanoseconds)</li> <li>Because of this slowness, DRAM is designed to use parallelism to increase the memory access throughput</li> </ul> Each time a DRAM location is accessed, a range of consecutive locations nearby are accessed. Once detected by the sensors, the data from all these locations can be transferred to the processor. <ul> <li>These consecutive locations assessed is known as DRAM bursts</li> <li>Making use of these consecutive and simultaneous accesses can allow us to improve memory access throughput</li> </ul> When a warp of threads is launched, all threads in the warp execute the same instruction at the same point in time. The hardware detects whether the threads are accessing consecutive memory locations. <ul> <li>Thus, the best access pattern is for threads in a warp to access consecutive data locations in global memory</li> <li>We say that the hardware coalesces the accesses into a consoliated access to consecutive DRAM locations</li> </ul> Recall that multidimensional data are linearized according to the row-major convention and stored into memory space. The read accesses are coalesced if consecutive threads in a warp are assigned to consecutive memory elements in this row major storage format. <h3 id="row-major"><a class="header" href="#row-major">Row Major</a></h3> Consider the following matrix multiplication code in which each thread computes one element of the output matrix <code>P</code> without tiling: <pre><code class="language-cpp">unsigned int row = blockIdx.y*blockDim.y + threadIdx.y; unsigned int col = blockIdx.x*blockDim.x + threadIdx.x; if (row < Width && col < Width) { float Pvalue = 0.0f; for (unsigned int k = 0; k < Width; ++k) { Pvalue += N[row*Width + k]*M[k*Width + col]; } P[row*Width + col] = Pvalue; } </code></pre> Notice that memory accesses to <code>M</code> are coalesced because: <ul> <li>The index of <code>M</code> is <code>k * Width + col</code></li> <li><code>k</code> and <code>Width</code> have the same values across all threads in a warp</li> <li><code>col = blockIdx.x * blockDim.x + threadIdx.x</code>, which means that consecutive threads (with incrementing <code>threadIdx.x</code>) will access consecutive elements</li> </ul> <h3 id="column-major"><a class="header" href="#column-major">Column Major</a></h3> Now suppose the matrix is stored in column major format instead. Now the code adjusts to the following: <pre><code class="language-cpp">unsigned int row = blockIdx.y*blockDim.y + threadIdx.y; unsigned int col = blockIdx.x*blockDim.x + threadIdx.x; if (row < Width && col < Width) { float Pvalue = 0.0f; for (unsigned int k = 0; k < Width; ++k) { Pvalue += N[row*Width + k]*M[col*Width + k]; } P[row*Width + col] = Pvalue; } </code></pre> Observe that the index to <code>M</code> changed from <code>k*Width + col</code> to <code>col*Width + k</code>. This is because in the column major format, elements that were previously in adjacent columns to each other are now in adjacent rows. Hence the indexing to <code>M</code> is swapped to reflect this. Now we can see that we have lost the coalesced access to <code>M</code>: <ul> <li>The indexing to <code>M</code> is now <code>col*Width + k</code></li> <li>Consecutive threads still have consecutive values of <code>col</code></li> <li>But <code>col</code> is now multiplied by <code>Width</code>, meaning that consecutive threads access values that are <code>Width</code> apart, which can be very far</li> </ul> <h3 id="corner-turning"><a class="header" href="#corner-turning">Corner Turning</a></h3> This is a method to create coalesced memory accesses even when the computation is not naturally amenable to it. The idea is to read from global memory into shared memory in a coalesced manner, then do the unfavourable access pattern from the cheap shared memory instead. Suppose we are in the tiled matrix multiplication example. First, we do a coalesced read from <code>M</code> into <code>Mds</code> for our tile by swapping the row and column. The elements can be arranged in <code>Mds</code> in either row or column major format, it does not matter. Then, we can do fast reads from <code>Mds</code> to do the matrix multiplication for our tile. <h2 id="hiding-memory-latency"><a class="header" href="#hiding-memory-latency">Hiding Memory Latency</a></h2> DRAM bursting is a form of parallel organization: multiple locations are accessed in the DRAM core array in parallel. There are two more ways of parallel organization: banks and channels. A single processor may have <code>8</code> channels. Each channel has a bus that accesses many banks. The bus transfers data from the banks back to the processor. The data transfer bandwidth of the bus is defined by its width and clock frequency. Modern DDR buses perform two data transfers per clock cycle. <ul> <li>e.g. A 64-bit DDR bus with clock frequency of 1 GHz has a bandwidth of <code>8B * 2 * 1 = 16GB/s</code>.</li> </ul> For each channel, we need many banks in order to fully utilize the data transfer bandwidth of the bus. Why? <ul> <li>Recall earlier that the data access latency is slow, because it takes time for cells to share their stored charge with the sensing amplifier</li> <li>Once the sensing amplifier has completed its work, the burst data is delivered through the bus (this part is relatively much faster)</li> <li>If we only have a single bank, there is a long latency waiting for the sensing amplifier while the bus sits idle</li> <li>If we have multiple buses, we can trigger multiple read accesses to different banks in close succession in an async manner</li> <li>When each bank has been sensed, the delivery of burst data can be tightly packed and fill up the idle time</li> </ul> In general, if the ratio of <code>cell array access latency / data transfer time = R</code>, then we need to have at least <code>R + 1</code> banks to fully utilize the data transfer bandwidth of the bus. <ul> <li>Note that we need more than <code>R</code> because of bank conflict. Since each bank can only serve one data access at a time (each bank can store multiple data points), if multiple threads need to access the same bank for data, there is a conflict. So having more banks help to reduce the probability of bank conflict</li> </ul> There is an important connection between occupancy and parallel access to DRAM data. As we saw earlier, maximizing occupancy ensures that we have enough threads on the SMs to hide long memory access latencies. In this section, we see another distinct advantage of maximizing occupancy - having many threads access different memory locations allows us to fully utilize the memory transfer bandwidth of the buses. That is, if the memory accesses are well distributed across channels and banks. Practically, the distribution of data across channels and banks is in an interleaved data distribution. Suppose we have a toy example: <ul> <li><code>4</code> channels</li> <li><code>2</code> banks per channel</li> <li>Each bank stores <code>2</code> consecutive floats</li> </ul> Now suppose we are doing a tiled matrix multiplication between two <code>4 x 4</code> matrices <code>M @ N</code>, where we set the tile width to <code>2</code>. The storage of <code>M</code> is like so (referring to the linearized index): <ul> <li>Channel 0 <ul> <li>Bank 0: <code>M[0], M[1]</code></li> <li>Bank 1: <code>M[8], M[9]</code></li> </ul> </li> <li>Channel 1 <ul> <li>Bank 0: <code>M[2], M[3]</code></li> <li>Bank 1: <code>M[10], M[11]</code></li> </ul> </li> <li>Channel 2 <ul> <li>Bank 0: <code>M[4], M[5]</code></li> <li>Bank 1: <code>M[12], M[13]</code></li> </ul> </li> <li>Channel 3 <ul> <li>Bank 0: <code>M[6], M[7]</code></li> <li>Bank 1: <code>M[14], M[15]</code></li> </ul> </li> </ul> As we can see, consecutive elements of our data arrays is first put into <code>bank 0</code> of all the channels. It then moves on to <code>bank 1</code> of all the channels. If there is more data, we will wrap back around to <code>bank 0</code> again. The breadth-first distribution across channels helps ensure that even when the amount of data is small, we have good distribution of data across channels and banks. Now for the matrix multiplication, we can verify that in phase 0 of the tiled matrix multiplication, these are the elements of <code>M</code> requested by each block: <ul> <li><code>Block 0,0</code>: <code>M[0], M[1], M[4], M[5]</code></li> <li><code>Block 0,1</code>: <code>M[0], M[1], M[4], M[5]</code></li> <li><code>Block 1,0</code>: <code>M[8], M[9], M[12], M[13]</code></li> <li><code>Block 1,1</code>: <code>M[8], M[9], M[12], M[13]</code></li> </ul> This is a bit tedious so will not elaborate on how these were derived. But basically, we can see that we access both banks of channels 0 and 2 for phase 1. So we are utilizing 50% of the channels. With multiplication of larger matrices, we can fully utilize all channels and buses. <h2 id="thread-coarsening"><a class="header" href="#thread-coarsening">Thread Coarsening</a></h2> So far, we have focused on kernels where work is parallelized across all threads at the finest granularity. For e.g. in the matrix multiplication example, each thread is assigned one element in the output matrix <code>P</code>. The advantage of parallelizing at the finest granularity is that it enhances transparent scalability - if the hardware has enough resources to run all the work in parallel, then we have maximized the parallelism. However, if the hardware does not have enough resources to run all the work at once, then a queue is formed up, and the parallel work becomes sequential. In this case, there is a price to be paid for such aggressive parallelism, such as: <ul> <li>redundant loading of the same data by different blocks</li> <li>redundant work</li> <li>syncrhonization overhead</li> </ul> When the threads are executed in parallel, this price is often well worth it. But if the hardware ends up serializing the work as a result of insufficient resources, this price may have been paid unnecessarily. It may thus be worthwhile to reduce the amount of parallelism by having each thread do more work, in exchange for reducing the price that we pay. One such technique is called thread coarsening. We can see a good example in the tiled matrix multiplication example. For calculating a given tile in the output matrix <code>P</code>, each thread block must load its own copy of the input tiles of matrix <code>M</code>. However, this means that the same tile of input matrix <code>M</code> will be loaded redundantly by different thread blocks and incurring access to global memory. More concretely, consider that we are computing rows of <code>P</code> from <code>P[i, :]</code> to <code>P[j, :]</code>. <ul> <li>One thread block may be computing <code>P[i, k1] to P[j, k2]</code></li> <li>One thread block may be computing <code>P[i, k2] to P[j, k3]</code></li> <li>But notice that to do this computation, they all have to access <code>M[i, k1] to M[j, k2]</code> in one of their phases</li> <li>It may make more sense for a single thread block to compute a few adjacent <code>P</code> tiles at once, and thus load this slice of matrix <code>M</code> only once and reduce redundant memory accesses</li> </ul> Here we can dive into the thread coarsening matrix multiplication kernel: <pre><code class="language-c">#define TILE_WIDTH 32 #define COARSE_FACTOR 4 __global__ void matrixMulKernel(float* M, float* N, float* P, int width) { __shared__ float Mds[TILE_WIDTH][TILE_WIDTH]; __shared__ float Nds[TILE_WIDTH][TILE_WIDTH]; int bx = blockIdx.x; int by = blockIdx.y; int tx = threadIdx.x; int ty = threadIdx.y; int row = by * TILE_WIDTH + ty; int colStart = bx * TILE_WIDTH * COARSE_FACTOR + tx; float Pvalue[COARSE_FACTOR]; for (int c = 0; c < COARSE_FACTOR; ++c) { Pvalue[c] = 0.0f; } for (int ph = 0; ph < width/TILE_WIDTH; ++ph) { Mds[ty][tx] = M[row*width + ph*TILE_WIDTH + tx]; for (int c = 0; c < COARSE_FACTOR; ++c) { int col = colStart + c * TILE_WIDTH; Nds[ty][tx] = N[(ph * TILE_WIDTH + ty) * width + col]; __syncthreads(); for (int k = 0; k < TILE_WIDTH; ++k) { Pvalue[c] += Mds[ty][k] * Nds[k][tx]; } __syncthreads(); } } for (int c = 0; c < COARSE_FACTOR; ++c) { int col = colStart + c * TILE_WIDTH; P[row * width + col] = Pvalue[c]; } } </code></pre> Let's unpack the code above. Notice that instead of each thread block processing one tile of the output <code>P</code> matrix, now one thread block will process <code>COARSE_FACTOR = 4</code> adjacent tiles of <code>P</code> <ul> <li>Therefore, each thread now processes <code>4</code> cells instead of <code>1</code> cell</li> </ul> First we identify <code>row</code> and <code>colStart</code>: <ul> <li><code>row</code> is the row of <code>P</code> that we are computing. This is as per normal</li> <li><code>colStart</code> is the start of the column that we are computing. Because each thread is now doing <code>4</code> cells, we multiply by <code>COARSE_FACTOR</code> to start further down. It is also the "start" because there are 4 cells to compute, and this variable indicates the first cell.</li> </ul> Initialize <code>Pvalue</code>: instead of one float, we have an array of <code>4</code> floats. Now the main loop: <ul> <li>The outer loop is over phases, which is as per previous. Each phase loads a new tile of <code>M</code> (by moving rightwards) and a new tile of <code>N</code> (by moving downwards) to compute the same <code>P</code> tile in a tiled manner.</li> <li>Notice that we load a tile of <code>M</code> for this block that will be reused across multiple loops of the coarsening loop over <code>c</code>. <ul> <li>The way <code>M</code> is loaded is as per normal - we locate the correct <code>row</code>, and the column is determined by <code>ph * TILE_WIDTH + tx</code></li> </ul> </li> <li>The coarsening loop is where things are interesting: <ul> <li>For a given value of <code>c</code>, we are working on one tile of <code>N</code></li> <li>Each iteration of the loop moves us in a stride of <code>TILE_WIDTH</code> to the right, as reflected by the way <code>col</code> is computed as <code>colStart + c * TILE_WIDTH</code></li> <li>e.g. a given thread may compute <code>col = 0, 32, 64, 96</code> instead of <code>col = 0, 1, 2, 3</code></li> <li>Once we have computed <code>col</code> to account for coarsening, the indexing into <code>N</code> is as per previous</li> </ul> </li> </ul> So we see that for the same <code>M</code> tile that is loaded, we reuse it <code>COARSE_FACTOR</code> times for different <code>N</code> tiles. Thread coarsening is a common technique that can led to substantial performance improvements. <h3 id="pitfalls"><a class="header" href="#pitfalls">Pitfalls</a></h3> There are some potential pitfalls for thread coarsening: <ul> <li>Not to apply optimization when it is unnecessary. Thread coarsening is only beneficial when there is a price paid for parallelization, such as redundant data loading. It will not help when there is no price paid.</li> <li>Apply too much coarsening such that hardware resources become underutilized. By applying coarsening, we reduce the amount of parallelism exposed and can lead to underutilisation.</li> <li>Third pitfall is that our coarsened kernel ends up using more resources which can hurt occupancy, for example, by using more registers per thread or more shared memory per thread block.</li> </ul> <h2 id="checklist-of-optimizations"><a class="header" href="#checklist-of-optimizations">Checklist of optimizations</a></h2> <ol> <li>Maximizing occupancy <ul> <li>Idea: more parallel memory accesses to hide DRAM latency</li> <li>Strategy: tune parameters like threads per block, shared memory per block, registers per thread to maximize occupancy</li> </ul> </li> <li>Enabling coalesced global memory accesses <ul> <li>Idea: less global memory traffic and better utilization of DRAM burst to efficiently read data</li> <li>Strategies: <ul> <li>Transfer from global to shared memory in a coalesced manner and perform uncoalesced accesses in shared memory</li> <li>Rearrange mapping of threads to data</li> <li>Rearrange layout of data</li> </ul> </li> </ul> </li> <li>Minimize control divergence <ul> <li>Idea: Higher SIMD efficiency and have less idle cores during SIMD execution</li> <li>Strategies: <ul> <li>Rearrange mapping of threads to work and / or data</li> <li>Rearrange layout of data</li> </ul> </li> </ul> </li> <li>Tiling of reused data <ul> <li>Idea: Less global memory traffic by reusing tiles</li> <li>Strategies: <ul> <li>Place reused data in shared memory or registers so that we only read from global memory once</li> </ul> </li> </ul> </li> <li>Privatisation <ul> <li>Idea: Less contention of writes to a shared copy of data and serialization of atomic updates</li> <li>Strategy: <ul> <li>Apply partial updates to a private copy of the data and then updataing the universal copy when done</li> </ul> </li> </ul> </li> <li>Thread coarsening <ul> <li>Idea: Less redundant work, divergene and global memory traffic</li> <li>Strategy: <ul> <li>Assign multiple units of parallelism to each thread to reduce the price of parallelism when unnecessarily incurred</li> </ul> </li> </ul> </li> </ol> <h2 id="exercises-4"><a class="header" href="#exercises-4">Exercises</a></h2> <blockquote> <ol> <li>Write a matrix multiplication kernel function that corresponds to the design illustrated in Fig. 6.4.</li> </ol> </blockquote> Fig 6.4 is the corner turning example. Suppose we have matrices <code>M x N = P</code>. For simplicity, we assume square matrices of <code>Width</code>. The access to <code>M</code> remains the same as normal tiled matrix multiplication, but we want to change the way <code>N</code> is loaded into shared memory, since the assumption is that <code>N</code> is stored in column major format. Doing a recap: each block computes a tile of <code>P</code>. Let <code>T</code> denotes the <code>TILE_WIDTH</code>. The rows and columns are determined by the current values of <code>by</code> and <code>bx</code> respectively: <ul> <li>The first row of the block is <code>by * T</code></li> <li>The first column of the block is <code>bx * T</code></li> <li><code>ty</code> and <code>tx</code> gives us the coordinate of the exact element within the block</li> </ul> Now let's focus on <code>N</code>. Here we need to distinguish between the logical indexing into <code>N</code> (call it <code>N_logical</code>) and the indexing into <code>N</code> based on how it is stored (column major, call it <code>N_stored</code>). We start with the logical view: <ul> <li>Each phase of the tiled matmul moves us downward on <code>N</code></li> <li>So we start at row <code>ph * T</code> and work downwards</li> <li>The columns are fixed as <code>bx * T</code> to <code>bx * (T+1) - 1</code></li> </ul> Since <code>N</code> is stored in column major, we need to swap the indexing when referring to <code>N_stored</code>. So: <ul> <li>We start at row <code>bx * T</code> (denoted <code>N_rowstart</code> in code)</li> <li>We start at column <code>ph * T</code> (denoted <code>N_colstart</code> in code)</li> </ul> Now we have found the correct coordinates of the tile within <code>N</code>. This is where it gets tricky: we want to read the elements of this tile in a coalesced way, such that incrementing <code>tx</code> by <code>1</code> reads the next element in the linearized array. <ul> <li>So the first element is <code>N_rowstart * Width</code> corresponding to <code>N_logical[ph*T, bx*T]</code></li> <li>Second element is <code>N_rowstart * Width + 1</code> corresponding to <code>N_logical[ph*T + 1, bx*T]</code> etc.</li> <li>This means that we will be reading in the transpose of this N tile, and we need to transpose it again when loading into <code>Nds</code> so that we get the correct tile when doing the final multiplication</li> </ul> Hopefully this analysis sheds some light on the code below. <pre><code class="language-c"># define TILE_WIDTH 16 __global__ void matrixMulKernel(float* M, float* N, float* P, int Width) { __shared__ float Mds[TILE_WIDTH][TILE_WIDTH]; __shared__ float Nds[TILE_WIDTH][TILE_WIDTH]; int bx = blockIdx.x; int by = blockIdx.y; int tx = threadIdx.x; int ty = threadIdx.y; // Identify P[Row, Col] to work on int Row = by * TILE_WIDTH + ty; int Col = bx * TILE_WIDTH + tx; float Pvalue = 0; for (int ph = 0; ph < Width/TILE_WIDTH; ++ph) { // Load M, N into shared memory Mds[ty][tx] = M[Row*Width + ph*TILE_WIDTH + tx]; // Note: rowstart, colstart refers to N_stored int N_rowstart = bx * TILE_WIDTH; int N_colstart = ph * TILE_WIDTH; Nds[tx][ty] = N[(N_rowstart + ty)*Width + N_colstart + tx]; __syncthreads(); for (int k = 0; k < TILE_WIDTH; ++k) { // We index Nds as per usual since it is now row major Pvalue += Mds[ty][k] * Nds[k][tx]; } __syncthreads(); } P[Row*Width + Col] = Pvalue; } </code></pre> <blockquote> <ol start="2"> <li>For tiled matrix multiplication, of the possible range of values for BLOCK_SIZE, for what values of BLOCK_SIZE will the kernel completely avoid uncoalesced accesses to global memory? (You need to consider only square blocks.)</li> </ol> </blockquote> Not too sure about this one. Since a warp has <code>32</code> threads, and we read <code>32</code> elements simultaneously, it is good if block_size is a multiple of <code>32</code>, so that we get coalesced reads of <code>32</code> elements at once. <blockquote> <ol start="3"> <li>Consider the following CUDA kernel:</li> </ol> </blockquote> <pre><code class="language-c">__global__ void foo_kernel(float* a, float* b, float* c, float* d, float* e) { unsigned int i = blockIdx.x*blockDim.x + threadIdx.x; __shared__ float a_s[256]; __shared__ float bc_s[4*256]; a_s[threadIdx.x] = a[i]; for(unsigned int j = 0; j < 4; ++j) { bc_s[j*256 + threadIdx.x] = b[j*blockDim.x*gridDim.x + i] + c[i*4 + j]; } __syncthreads(); d[i + 8] = a_s[threadIdx.x]; e[i*8] = bc_s[threadIdx.x*4]; } </code></pre> <blockquote> For each of the following memory accesses, specify whether they are coalesced or uncoalesced or coalescing is not applicable: </blockquote> <blockquote> a. The access to array a of line 05 </blockquote> Coalesced because increasing <code>tx</code> by <code>1</code> takes the next element of <code>a</code> <blockquote> b. The access to array a_s of line 05 </blockquote> Not applicable, <code>a_s</code> is shared memory. <blockquote> c. The access to array b of line 07 </blockquote> Coalesced because increasing <code>tx</code> by <code>1</code> will take the next element of <code>b</code> <blockquote> d. The access to array c of line 07 </blockquote> Not coalesced because increasing <code>tx</code> by <code>1</code> will increase the index into <code>c</code> by <code>4</code>. <blockquote> e. The access to array bc_s of line 07 </blockquote> Not applicable, <code>bc_s</code> is shared memory. <blockquote> f. The access to array a_s of line 10 </blockquote> Not applicable, <code>a_s</code> is shared memory. <blockquote> g. The access to array d of line 10 </blockquote> Coalesced. <blockquote> h. The access to array bc_s of line 11 </blockquote> Not applicable. <blockquote> i. The access to array e of line 11 </blockquote> Not coalesced, increasing <code>tx</code> will increment access to <code>e</code> by <code>8</code>. <blockquote> <ol start="4"> <li>What is the floating point to global memory access ratio (in OP/B) of each of the following matrix-matrix multiplication kernels?</li> </ol> </blockquote> <blockquote> a. The simple kernel described in Chapter 3, Multidimensional Grids and Data, without any optimizations applied. </blockquote> In the simple matmul, each thread accesses <code>Width</code> number of elements of <code>M</code> and <code>N</code> respectively, and performs a <code>sum</code> and <code>multiplication</code> for each loop. <ul> <li><code>2 x Width</code> number of global memory accesses</li> <li><code>2 x Width</code> number of OPs</li> <li><code>OP/B = 2 x Width / (2 x Width x 4) = 1/4</code> (because <code>4 bytes</code> per float)</li> </ul> <blockquote> b. The kernel described in Chapter 5, Memory Architecture and Data Locality, with shared memory tiling applied using a tile size of 32 x 32. </blockquote> For tiled matmul (see <a href="pmpp/ch05.html#a-tiled-matmul-kernel">tiled matmul</a>), each thread still does <code>2 x Width</code> number of OPs. But the number of global memory accesses is now <code>2 x Width / TILE_WIDTH</code>. <ul> <li><code>OP / B = 2 x Width / (2 x Width x 4 / 32) = 8 OP/B</code></li> </ul> <blockquote> c. The kernel described in this chapter with shared memory tiling applied using a tile size of 32 x 32 and thread coarsening applied using a coarsening factor of 4. </blockquote> <a href="pmpp/ch06.html#thread-coarsening">Thread coarsening</a> means that we reuse the same <code>M</code> tile for computing <code>4</code> output tiles: <ul> <li><code>Width / TILE_WIDTH x COARSE_FACTOR</code> number of accesses for <code>N</code></li> <li><code>Width / TILE_WIDTH</code> number of accesses for <code>M</code></li> <li><code>2 x Width x COARSE_FACTOR</code> number of FLOPs</li> <li>So numerator is <code>2 x Width</code></li> <li>Denominator is <code>4 x Width / (TILE_WIDTH x COARSE_FACTOR) x Width / (TILE_WIDTH)</code></li> <li><code>OP / B = 2/4 / (1/128 + 1/32) = 12.8 OP/B</code></li> </ul> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="mackay---info-theory"><a class="header" href="#mackay---info-theory">Mackay - info theory</a></h1> Notes on David Mackay - Information Theory, Inference and Learning Algorithms <a href="https://www.inference.org.uk/mackay/itila/book.html">link</a> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="preface"><a class="header" href="#preface">Preface</a></h1> The book provides some math as preface to understand the first chapter. <h2 id="the-binomial-distribution"><a class="header" href="#the-binomial-distribution">The Binomial Distribution</a></h2> <h3 id="example-11---binomial"><a class="header" href="#example-11---binomial">Example 1.1 - Binomial</a></h3> <blockquote> A coin has probability f of coming up heads. The coin is tossed N times. What is the probability distribution of the number of heads, r? What are the mean and variance of r? </blockquote> The number of heads has a binomial distribution. Each instance has probability as such, and there are (rN) ways that it can occur. P(r∣f,N)=(rN)fr(1−f)N−r The mean is: E[r]≡r=0∑NP(r∣f,N)r The variance is: var[r]≡E[(r−E[r])2]=E[rr]−(E[r])2=r=0∑NP(r∣f,N)r2−(E[r])2 Computing these sums are cumbersome. It is easier to observe that r (the number of heads) is the sum of N independent random variables - each random variable is whether the ith toss results in a head. In general, for any random variables x and y, if x and y are independent: E[x+y]var[x+y]=E[x]+E[y]=var(x)+var(y) So the mean: E[r]=N⋅[f×1+(1−f)×0]=Nf And the variance (note that we subtract f2 as the expectation of a single trial is f): var[r]=N⋅[f×12+(1−f)×02−f2]=Nf(1−f) <h2 id="stirlings-approximation"><a class="header" href="#stirlings-approximation">Stirling's Approximation</a></h2> We will use the poisson distribution and the gaussian approximation to derive stirling's approximation for the factorial function. Start with a poisson distribution with mean λ: P(r∣λ)=e−λr!λr r∈{0,1,2,...} For large λ, the poisson distribution is well approximated in the vicinity of r≈λ by a gaussian distribution with mean λ and variance λ: e−λr!λr≈2πλ<svg xmlns="http://www.w3.org/2000/svg" width='400em' height='1.08em' viewBox='0 0 400000 1080' preserveAspectRatio='xMinYMin slice'><path d='M95,702 c-2.7,0,-7.17,-2.7,-13.5,-8c-5.8,-5.3,-9.5,-10,-9.5,-14 c0,-2,0.3,-3.3,1,-4c1.3,-2.7,23.83,-20.7,67.5,-54 c44.2,-33.3,65.8,-50.3,66.5,-51c1.3,-1.3,3,-2,5,-2c4.7,0,8.7,3.3,12,10 s173,378,173,378c0.7,0,35.3,-71,104,-213c68.7,-142,137.5,-285,206.5,-429 c69,-144,104.5,-217.7,106.5,-221 l0 -0 c5.3,-9.3,12,-14,20,-14 H400000v40H845.2724 s-225.272,467,-225.272,467s-235,486,-235,486c-2.7,4.7,-9,7,-19,7 c-6,0,-10,-1,-12,-3s-194,-422,-194,-422s-65,47,-65,47z M834 80h400000v40h-400000z'/></svg>1e−2λ(r−λ)2 Plugging r=λ, we get: e−λλ!λλλ!≈2πλ<svg xmlns="http://www.w3.org/2000/svg" width='400em' height='1.08em' viewBox='0 0 400000 1080' preserveAspectRatio='xMinYMin slice'><path d='M95,702 c-2.7,0,-7.17,-2.7,-13.5,-8c-5.8,-5.3,-9.5,-10,-9.5,-14 c0,-2,0.3,-3.3,1,-4c1.3,-2.7,23.83,-20.7,67.5,-54 c44.2,-33.3,65.8,-50.3,66.5,-51c1.3,-1.3,3,-2,5,-2c4.7,0,8.7,3.3,12,10 s173,378,173,378c0.7,0,35.3,-71,104,-213c68.7,-142,137.5,-285,206.5,-429 c69,-144,104.5,-217.7,106.5,-221 l0 -0 c5.3,-9.3,12,-14,20,-14 H400000v40H845.2724 s-225.272,467,-225.272,467s-235,486,-235,486c-2.7,4.7,-9,7,-19,7 c-6,0,-10,-1,-12,-3s-194,-422,-194,-422s-65,47,-65,47z M834 80h400000v40h-400000z'/></svg>1≈λλe−λ2πλ<svg xmlns="http://www.w3.org/2000/svg" width='400em' height='1.08em' viewBox='0 0 400000 1080' preserveAspectRatio='xMinYMin slice'><path d='M95,702 c-2.7,0,-7.17,-2.7,-13.5,-8c-5.8,-5.3,-9.5,-10,-9.5,-14 c0,-2,0.3,-3.3,1,-4c1.3,-2.7,23.83,-20.7,67.5,-54 c44.2,-33.3,65.8,-50.3,66.5,-51c1.3,-1.3,3,-2,5,-2c4.7,0,8.7,3.3,12,10 s173,378,173,378c0.7,0,35.3,-71,104,-213c68.7,-142,137.5,-285,206.5,-429 c69,-144,104.5,-217.7,106.5,-221 l0 -0 c5.3,-9.3,12,-14,20,-14 H400000v40H845.2724 s-225.272,467,-225.272,467s-235,486,-235,486c-2.7,4.7,-9,7,-19,7 c-6,0,-10,-1,-12,-3s-194,-422,-194,-422s-65,47,-65,47z M834 80h400000v40h-400000z'/></svg> Applying ln, we get stirling's approximation for the factorial function: x!lnx!≈xxe−x2πx<svg xmlns="http://www.w3.org/2000/svg" width='400em' height='1.08em' viewBox='0 0 400000 1080' preserveAspectRatio='xMinYMin slice'><path d='M95,702 c-2.7,0,-7.17,-2.7,-13.5,-8c-5.8,-5.3,-9.5,-10,-9.5,-14 c0,-2,0.3,-3.3,1,-4c1.3,-2.7,23.83,-20.7,67.5,-54 c44.2,-33.3,65.8,-50.3,66.5,-51c1.3,-1.3,3,-2,5,-2c4.7,0,8.7,3.3,12,10 s173,378,173,378c0.7,0,35.3,-71,104,-213c68.7,-142,137.5,-285,206.5,-429 c69,-144,104.5,-217.7,106.5,-221 l0 -0 c5.3,-9.3,12,-14,20,-14 H400000v40H845.2724 s-225.272,467,-225.272,467s-235,486,-235,486c-2.7,4.7,-9,7,-19,7 c-6,0,-10,-1,-12,-3s-194,-422,-194,-422s-65,47,-65,47z M834 80h400000v40h-400000z'/></svg>≈xlnx−x+21ln2πx Now we apply stirling's approximation to ln(rN) and re-organize: ln(rN)≡ln(N−r)!r!N!≈NlnN−(N−r)ln(N−r)−rlnr≈NlnN−rN−rlnN−r1+rlnr1≈(N−r)lnN−rN+rlnrN Note that: <ul> <li>We are using the approximation of lnx!≈xlnx−x</li> <li>In the second line, we replace all the factorials with the approximation, and notice that the −x part of the approximation will cancel out between all the terms</li> <li>In the third line, we group terms together and use the trick lnx=−lnx1 to flip signs</li> <li>In the last line, we artificially add and subtract rlnN to create the final form</li> </ul> Now, the whole point of this derivation is to write the approximation in a form that resembles binary entropy. We now define that concept here. <h3 id="binary-entropy"><a class="header" href="#binary-entropy">Binary Entropy</a></h3> The binary entropy function is (log denotes logarithm with base 2): H2(x)≡xlogx1+(1−x)log1−x1 <blockquote> Note that the binary entropy function is just the special case of the general shannon entropy when we have only two possible outcomes, assuming x represents a probability. </blockquote> Let us rewrite log(rN) using the binary entropy function: log(RN)≈NH2(r/N)=N(NrlogrN+NN−rlogN−rN)=(N−r)logN−rN+rlogrN Or equivalently, (rN)≈2NH2(r/N) So we see that the binary entropy function nicely describes the combinatorial explosion of the binomial function. We can observe that since the binary entropy function is maximized when Nr=1/2, the binomial combination is maximized likewise. <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="chapter-1-introduction-to-information-theory"><a class="header" href="#chapter-1-introduction-to-information-theory">Chapter 1: Introduction to Information Theory</a></h1> The first problem we tackle is how to communicate perfectly over imperfect communication channels. The running example is this. We have some message which is a sequence of bits. We want to write this data to a noisy hard disk that transmits each bit correctly with probability 1−f and correctly with probability f. This is called the binary symmetric channel (symmetric as the probablity of flipping from 0 to 1 is the same as 1 to 0). Suppose f=0.1. This is clearly unacceptable. We want a useful disk drive to flip no bits in say 10 years writing 1GB a day, which means we want error probability around 10−15 or smaller. We solve this problem by introducing communication systems to detect and correct the errors. The flow is such: <ul> <li>Start with message s</li> <li>s gets encoded into t, which is some intermediate representation we design</li> <li>Some noise n gets added to t which represents noise from the imperfect hard disk</li> <li>This results in r=t+n which is the received vector <ul> <li>The addition is in modulo 2 arithmetic, i.e. 0+1=1 and 1+1=0</li> </ul> </li> <li>Finally we decode using some algorithm to get s^ from r</li> </ul> <h2 id="the-r3-repetition-code"><a class="header" href="#the-r3-repetition-code">The R3 Repetition Code</a></h2> A simple communication system is to repeat each bit 3 times. After noise is added, we decode each bit in the received vector by taking the majority vote. The idea is that when f is small, the probability of multiple bits failing independently becomes small. So the majority vote will correct some errors. <blockquote> Theorem. The majority vote decoding is optimal (i.e. has the lowest probability of error). </blockquote> Proof. Since each bit is independent, we just need to consider the optimal decision for one bit. Suppose we received 3 bits r1,r2,r3. The objective is to recover the original bit s, which could have been 0 or 1. Thus, we want: s^=s∈{0,1}argmaxP(S=s∣r1,r2,r3)=sargmaxP(r1,r2,r3)P(r1,r2,r3∣s)P(s)=sargmaxP(r1,r2,r3∣s) Note that: <ul> <li>In the first line, S is the random variable representing the unknown true bit</li> <li>s is a candidate value that we are choosing (some notation overloading as s is also used to represent the true bit above)</li> <li>P(S=s) is treated as a constant, and P(r1,r2,r3) does not affect choice of s</li> </ul> Now, if we denote d as the number of bit flips between r1,r2,r3 and s, we have: P(r1,r2,r3∣s)=(1−f)3−dfd Notice that we are computing the probability for this specific received sequence r1,r2,r3, hence we do not need to multiply by the number of combinations. With this formulation, we see that P(r1,r2,r3∣s) is maximized when the choice of s minimizes the number of bit flips (if f<0.5). This is because a scenario with few bit flips is more likely than a scenario with many bit flips. Hence, the majority vote is optimal as it minimizes the number of bit flips. <blockquote> Exercise 1.2. Show that the error probability is reduced by the user of R3 by computing the error probability of this code for a binary symmetric channel with noise level f. </blockquote> Solution. The naive error probability is: P(r=s)=f With R3, the error probability is: P(rmajority=s)=P(2 or 3 bit flips)=f2(1−f)×3+f3=3f2−2f3 Now we want to show that when f<0.5, R3 has a better error probability: f3f3f−2f23f2−2f3<0.5<1.5<1<f Hence we showed that R3 has a better error probability. Even though R3 has a better error probability, the problem is that our rate of information transfer has fallen by a factor of 3. To improve our error probability, we can continue to increase repetition at the cost of decreased rate. <blockquote> Exercise 1.3. Find the probability of error of RN, the repetition code with N repetitions for odd N </blockquote> P(Error of RN)=P(At least 2N+1 bits flipped)=n=(N+1)/2∑N(nN)fn(1−f)N−n <blockquote> Assuming f=0.1, which terms in this sum is biggest? How much bigger than the second largest term? </blockquote> The largest term is ((N+1)/2N)f(N+1)/2(1−f)(N−1)/2 The second largest term is around 0.1 times smaller, so the largest term dominates. <blockquote> Use stirling's approximation to approximate the largest term and find the probability of error. </blockquote> Using stirling's approximation: ((N+1)/2N)≈2N⋅H2((N+1)/2/N)≈2N⋅H2(1/2)=2N Hence we have probability of error as: Perror≈2N⋅fN/2⋅(1−f)N/2=(4f(1−f))N/2 To get error probability less than 10−15, we write out the inequality and manipulate to find the lower bound for N. <h2 id="hamming-7-4-code"><a class="header" href="#hamming-7-4-code">Hamming (7, 4) code</a></h2> Thus far, we have been trying to encode each bit independently. What if we encode blocks of bits together? Can we get more efficient codes? A block code converts a sequence of source bits s, of length K, into a transmitted sequence t of length N bits. <ul> <li>We add some redundancy by making N>K. The extra N−K bits (called parity-check bits) are used to store some redundant information of the original K bits.</li> <li>This is usually implemented as some linear function of the original K bits</li> <li>In the (7,4) Hamming Code, we transmit N=7 bits for every K=4 source bits</li> </ul> The encoding can be shown with an example. Suppose s=0011. <ul> <li>The first 4 bits of transmitted t=0011, i.e. just copy the source sequence <ul> <li>t1,t2,t3,t4=s1,s2,s3,s4</li> </ul> </li> <li>The next 3 bits are parity-check bits <ul> <li>t5 is set such that s1+s2+s3+t5 is even</li> <li>t6 is set such that s2+s3+s4+t6 is even</li> <li>t7 is set such that s1+s3+s4+t7 is even</li> </ul> </li> <li>For the example of s=0011: <ul> <li>t=0011100</li> </ul> </li> </ul> We can see that the Hamming code is a linear code, since the encoding can be written compactly in terms of a matrix-vector multiplication (using modulo-2 arithmetic). Specifically, the transmitted code-vector t may be obtained from source-vector s using a matrix multiplication: t=GTs Where G is the generator matrix of the code: GT=<svg xmlns="http://www.w3.org/2000/svg" width='0.667em' height='8.400em' viewBox='0 0 667 8400'><path d='M403 1759 V84 H666 V0 H319 V1759 v4800 v1759 h347 v-84 H403z M403 1759 V0 H319 V1759 v4800 v1759 h84z'/></svg>1000101010011000101110001011<svg xmlns="http://www.w3.org/2000/svg" width='0.667em' height='8.400em' viewBox='0 0 667 8400'><path d='M347 1759 V0 H0 V84 H263 V1759 v4800 v1759 H0 v84 H347z M347 1759 V0 H263 V1759 v4800 v1759 h84z'/></svg> The columns of the generator matrix may be viewed as defining four basis vectors in a seven dimensional binary space. The sixteen codewords are obtained by taking all possible linear combinations of these vectors (in modulo-2 arithmetic). This is a useful perspective: <ul> <li>With 7 dimensions, we have <code>128</code> possible vectors, out of which only <code>2^4 = 16</code> are valid. This means that if we get an invalid vector, we know that it needs to be corrected.</li> <li>The <code>16</code> valid vectors are well spread-out in the space (every pair of codewords differs in at least 3 positions), enabling correction</li> </ul> <h3 id="decoding-hamming-code"><a class="header" href="#decoding-hamming-code">Decoding Hamming Code</a></h3> Decoding the hamming code follows the same logic as before. Assuming that the channel is a binary symmetric channel and all source vectors are equiprobable, we want to choose a source vector s whose encoding t(s) differs from received r in as few bits as possible. Recall that the theorem we proved earlier shows that we want to maximize p(r1,...,r7∣s). Naively, we can enumerate all <code>16</code> valid hamming codes (each of length <code>7</code>), and then compare t(s) against each and find the closest. This is inefficient as our codes get longer. The better way is syndrome decoding. A syndrome is defined as the pattern of violations of the parity bits. In other words, it is the vector of "unhappy" parity bits. <blockquote> Syndome example. Suppose we transmit <code>t = 1000101</code> and noise flips the second bit, giving us <code>r = 1100101</code>. The syndrome would be <code>z = 110</code>, because the first two parity checks are unhappy and the third is happy. </blockquote> The syndrome decoding task is thus to find the unique bit that lies inside all the unhappy circles and outside the happy circles. If we can find such a bit, flipping it would account for the observed syndrome. <h3 id="matrix-version-of-hamming-decoding"><a class="header" href="#matrix-version-of-hamming-decoding">Matrix Version of Hamming Decoding</a></h3> We can describe the decoding operation in terms of matrices. Let us define: P=<svg xmlns="http://www.w3.org/2000/svg" width='0.667em' height='3.600em' viewBox='0 0 667 3600'><path d='M403 1759 V84 H666 V0 H319 V1759 v0 v1759 h347 v-84 H403z M403 1759 V0 H319 V1759 v0 v1759 h84z'/></svg>101110111011<svg xmlns="http://www.w3.org/2000/svg" width='0.667em' height='3.600em' viewBox='0 0 667 3600'><path d='M347 1759 V0 H0 V84 H263 V1759 v0 v1759 H0 v84 H347z M347 1759 V0 H263 V1759 v0 v1759 h84z'/></svg>∈R3×4 Let H=[−PI3]. Then the syndrome vector is: z=Hr Since the −P part computes the expected parity bit and subtracts it from the actual received parity bits (from the identity part). But because in modulo-2 arithmetic, −1≡1, so: H=[PI3]=<svg xmlns="http://www.w3.org/2000/svg" width='0.667em' height='3.600em' viewBox='0 0 667 3600'><path d='M403 1759 V84 H666 V0 H319 V1759 v0 v1759 h347 v-84 H403z M403 1759 V0 H319 V1759 v0 v1759 h84z'/></svg>101110111011100010001<svg xmlns="http://www.w3.org/2000/svg" width='0.667em' height='3.600em' viewBox='0 0 667 3600'><path d='M347 1759 V0 H0 V84 H263 V1759 v0 v1759 H0 v84 H347z M347 1759 V0 H263 V1759 v0 v1759 h84z'/></svg> It should be clear that for valid transmitted vectors of the form t=GTs, Ht will be the zero vector / syndrome. <blockquote> Exercise 1.4. Prove that this is so. </blockquote> We want to show that HGTs is the zero vector for all s. Consider HGT. We have: HGT=[PI3][I4P] For compatible matrix blocks (i.e. P⋅I4 is defined etc.), we have: HGT=P+P=0∈R3×4 Hence HGTs will always be the zero vector. <hr /> We close this section of hamming decoding by noting that the received vector is obtained by: r=GTs+n And the syndrome vector is: z=Hr=HGTs+Hn=Hn Hence, the syndrome decoding problem is to find the most probably noise vector that satisfies z=Hn. Such a decoding algorithm is called maximum likelihood decoder. <blockquote> Exercise 1.5 Refer to the (7,4) Hamming code. Decode these received strings: <code>r = 1101011</code> </blockquote> We have <code>z = 011</code>, so flip r4, s^=<code>1100</code> <blockquote> <code>r = 0110110</code> </blockquote> We have <code>z = 111</code>, so flip r3, s^=<code>0100</code> <blockquote> <code>r = 0100111</code> </blockquote> We have <code>z = 001</code>, so flip r7, s^=<code>0100</code> <blockquote> <code>r = 1111111</code> We have <code>z = 000</code>, so flip nothing, s^=<code>1111</code> </blockquote> <blockquote> Exercise 1.6. Calculate the probability of block error pB of the 7, 4 hamming code as a function of the noise level f and show that to leading order it goes as 21f2. </blockquote> The block error is the probability that one or more of the decoded bits in one block fail to match the corresponding source bits. Assuming that whenever 2 or more bits are flipped in a block of 7 bits, we get a block decoding error, we can derive this as a binomial distribution. pB=P(At least 2 bits flipped)=1−P(0 bits flipped)−P(1 bit flipped)=1−(1−f)7−7f(1−f)6 Recall Taylor's expansion: (1−f)n=k=0∑n(kn)(−f)k So for pB: pB=1−(1−f)7−7f(1−f)6=1−(1−7f+21f2+...)−7f(1−6f+...)=21f2+O(f3) The leading term is 21f2. <blockquote> Show that to leading order the probability of bit error pb goes as 9f2. </blockquote> <div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="the-kernel-hacker-roadmap"><a class="header" href="#the-kernel-hacker-roadmap">The "Kernel Hacker" Roadmap</a></h1> Roadmap for learning and understanding LLM internals, unsloth optimizations, customized triton kernels etc. <h2 id="phase-1-the-new-tooling-triton--block-pointers"><a class="header" href="#phase-1-the-new-tooling-triton--block-pointers">Phase 1: The New Tooling (Triton & Block Pointers)</a></h2> <blockquote> Shift from CUDA C++ mental models to Triton's "Block-Centric" Python. </blockquote> Learn how triton abstracts tiling and memory coalescing etc. <ul> <li>The Concept: In CUDA, you calculate <code>idx</code> per thread. In Triton, you calculate pointers for a whole block of data.</li> <li>Key Mechanic: <code>tl.make_block_ptr</code>. This is the modern way to handle 2D tiling without manual boundary checks inside the hot loop.</li> </ul> <h3 id="-action-items"><a class="header" href="#-action-items">📝 Action Items</a></h3> <ol> <li>Tutorial: Complete <a href="https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html">Triton Tutorial 03 (Matrix Multiplication)</a>.</li> <li>Challenge: Rewrite the MatMul tutorial using <code>make_block_ptr</code> (if the tutorial uses the old offset arithmetic).</li> <li>Deep Dive: Read the docs for <code>tl.load(boundary_check=...)</code>. Understand how Triton handles padding automatically when your tile goes off the edge of the tensor.</li> </ol> <hr /> <h2 id="phase-2-operator-fusion-the-liger-pattern"><a class="header" href="#phase-2-operator-fusion-the-liger-pattern">Phase 2: Operator Fusion (The "Liger" Pattern)</a></h2> <blockquote> Learn to fuse memory-bound operations to reduce HBM reads/writes. </blockquote> This is the easiest area to make immediate gains (e.g., Unsloth's 30% speedup on non-attention layers). <ul> <li>The Target: <code>RMSNorm</code>, <code>CrossEntropyLoss</code>, <code>GeLU</code>.</li> <li>The "Hack": Instead of <code>Read -> Softmax -> Write -> Read -> Log -> Write</code>, you keep data in SRAM and do <code>Read -> Softmax -> Log -> Write</code>.</li> <li>Mathematical Hurdle: "Online Softmax". You cannot see the whole row at once to calculate the sum for the denominator. You must calculate running max/sum statistics as you load tiles.</li> </ul> <h3 id="-action-items-1"><a class="header" href="#-action-items-1">📝 Action Items</a></h3> <ol> <li>Study: Clone <a href="https://github.com/linkedin/Liger-Kernel">Liger-Kernel</a>. Read <code>liger_kernel/ops/rms_norm.py</code>. It is cleaner than Unsloth’s code.</li> <li>Code: Implement a Fused RMSNorm in Triton. <ul> <li>Steps: Load vector → Cast to FP32 → Square → Mean → Sqrt → Cast back → Store.</li> </ul> </li> <li>Verify: Validate your kernel output against <code>torch.nn.RMSNorm</code>.</li> </ol> <hr /> <h2 id="phase-3-the-engine-attention--flash"><a class="header" href="#phase-3-the-engine-attention--flash">Phase 3: The Engine (Attention & Flash)</a></h2> <blockquote> Mastering the inverted loop and memory hierarchy. </blockquote> <ul> <li>The Core Insight: Standard attention materializes an N×N matrix. Flash Attention never materializes it.</li> <li>The Algorithm (Inverted Loop): <ol> <li>Load a block of Query (Q) into SRAM.</li> <li>Loop over blocks of Key/Value (K,V) from HBM.</li> <li>Compute scores and update the output accumulator using online softmax logic.</li> </ol> </li> <li>FA2 vs FA3: <ul> <li>FA2: Parallelizes over sequence length; uses standard Tensor Cores.</li> <li>FA3 (Hopper H100 only): Uses Warp Specialization (Producer warps load data, Consumer warps do math) and asynchronous TMA (Tensor Memory Accelerator) copies.</li> </ul> </li> </ul> <h3 id="-action-items-2"><a class="header" href="#-action-items-2">📝 Action Items</a></h3> <ol> <li>Read: <a href="https://arxiv.org/abs/2307.08691">FlashAttention-2 Paper</a> (Focus on the Algorithm pseudocode, not the proofs).</li> <li>Tutorial: <a href="https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html">Triton Tutorial 06 (Fused Attention)</a>. <ul> <li>Crucial: Annotate the code where <code>m_i</code> (running max) and <code>l_i</code> (running sum) are updated. This is the heart of the algorithm.</li> </ul> </li> </ol> <hr /> <h2 id="phase-4-llm-specifics-rope-packing-precision"><a class="header" href="#phase-4-llm-specifics-rope-packing-precision">Phase 4: LLM Specifics (RoPE, Packing, Precision)</a></h2> <blockquote> Implementing the mechanisms that make modern LLMs work. </blockquote> <h3 id="a-rotary-positional-embeddings-rope"><a class="header" href="#a-rotary-positional-embeddings-rope">A. Rotary Positional Embeddings (RoPE)</a></h3> <ul> <li>The Insight: RoPE is a rotation in a complex plane. It is memory-bound.</li> <li>Implementation: You don't use matrix multiplication. You apply element-wise rotation: xnew=xreal⋅cosθ−ximag⋅sinθ</li> <li>Optimization: In optimized kernels (Unsloth/FA), RoPE is applied inside the attention kernel after loading Q/K to SRAM, but before the dot product.</li> </ul> <h3 id="b-sequence-packing-effective-batching"><a class="header" href="#b-sequence-packing-effective-batching">B. Sequence Packing (Effective Batching)</a></h3> <ul> <li>The Problem: Padding tokens waste compute.</li> <li>The Solution: Flatten batch into 1D stream <code>[Total_Tokens, Dim]</code>.</li> <li>The Implementation: <ul> <li><code>cu_seqlens</code>: A tensor <code>[0, 100, 550, ...]</code> defining document boundaries.</li> <li>Block Diagonal Masking: Your kernel uses <code>cu_seqlens</code> to know when to stop attending (i.e., don't let Doc A attend to Doc B).</li> <li>Exercise: Look at <code>flash_attn_varlen_func</code> in the Flash Attention repo to see how <code>cu_seqlens</code> is passed.</li> </ul> </li> </ul> <hr /> <h2 id="phase-5-the-hacker-skill-integration--monkey-patching"><a class="header" href="#phase-5-the-hacker-skill-integration--monkey-patching">Phase 5: The "Hacker" Skill (Integration & Monkey Patching)</a></h2> <blockquote> How Unsloth actually works. </blockquote> Unsloth does not rewrite the model source code. It modifies the Python objects in memory at runtime. <h3 id="-action-items-3"><a class="header" href="#-action-items-3">📝 Action Items</a></h3> <ol> <li>The Pattern: Write a script to "Patch" a model. <pre><code class="language-python">import torch from transformers import AutoModelForCausalLM # 1. Your Custom Wrapper def my_fast_forward(self, x, ...): return my_triton_kernel(x) # 2. The Patch model = AutoModelForCausalLM.from_pretrained("llama-3-8b") # Replace the method on the class or the instance model.model.layers[0].self_attn.forward = my_fast_forward.__get__( model.model.layers[0].self_attn, type(model.model.layers[0].self_attn) ) </code></pre> </li> <li>Flex Attention: If writing raw kernels is too hard, learn <code>torch.nn.attention.flex_attention</code>. <ul> <li>Write a Python <code>score_mod</code> that implements a custom mask (e.g., "Sliding Window") and compile it.</li> </ul> </li> </ol> <hr /> <h2 id="-reference-library"><a class="header" href="#-reference-library">📚 Reference Library</a></h2> <div class="table-wrapper"><table><thead><tr><th style="text-align: left">Topic</th><th style="text-align: left">Resource</th><th style="text-align: left">Why use it?</th></tr></thead><tbody> <tr><td style="text-align: left">Foundations</td><td style="text-align: left">PMPP Book</td><td style="text-align: left">It explains the "Why" (Memory hierarchy).</td></tr> <tr><td style="text-align: left">Fusion Examples</td><td style="text-align: left"><a href="https://github.com/linkedin/Liger-Kernel">Liger-Kernel</a></td><td style="text-align: left">Readable, production-grade Triton kernels for Norms/Loss.</td></tr> <tr><td style="text-align: left">Attention Code</td><td style="text-align: left"><a href="https://triton-lang.org/main/getting-started/tutorials/index.html">Triton Tutorials</a></td><td style="text-align: left">The "Hello World" for GPU programming.</td></tr> <tr><td style="text-align: left">Architecture</td><td style="text-align: left"><a href="https://github.com/Dao-AILab/flash-attention">Flash Attention Repo</a></td><td style="text-align: left">The "Source of Truth" for how packing (<code>varlen</code>) is handled.</td></tr> <tr><td style="text-align: left">Hacking</td><td style="text-align: left"><a href="https://github.com/unslothai/unsloth">Unsloth Source</a></td><td style="text-align: left">Study <code>models/llama.py</code> to see how they inject kernels.</td></tr> </tbody></table> </div><div style="break-before: page; page-break-before: always;"></div><link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.4/dist/katex.min.css"> <h1 id="information-theory-roadmap"><a class="header" href="#information-theory-roadmap">Information Theory Roadmap</a></h1> A roadmap for learning information theory and how it applies to model compression and hallucinations. <h2 id="week-1-foundations-of-information-theory"><a class="header" href="#week-1-foundations-of-information-theory">Week 1: Foundations of Information Theory</a></h2> <ul> <li>Mackay Chapters 1 (Intro), 2 (Probability / Entropy), 4 (Source coding theorem)</li> <li>Check that minimizing cross entropy loss is identical to minimizing KL divergence</li> </ul> <h2 id="week-2-practical-compression-with-arithmetic-coding"><a class="header" href="#week-2-practical-compression-with-arithmetic-coding">Week 2: Practical Compression with Arithmetic Coding</a></h2> <ul> <li>Mackay Chapter 6 (stream codes)</li> <li>Implement a basic arithmetic coder to compress a string given a fixed probability table</li> </ul> <h2 id="week-3-bayesian-model-comparison"><a class="header" href="#week-3-bayesian-model-comparison">Week 3: Bayesian Model Comparison</a></h2> <ul> <li>Mackay Chapter 28 (model comparison)</li> </ul> <h2 id="week-4-concentration-inequalities"><a class="header" href="#week-4-concentration-inequalities">Week 4: Concentration Inequalities</a></h2> <ul> <li>High dimensional probability (Vershynin) Chapters 2-3.</li> </ul> <h2 id="week-5-universal-coding-and-mdl"><a class="header" href="#week-5-universal-coding-and-mdl">Week 5: Universal Coding and MDL</a></h2> <ul> <li>Cover and Thomas Chapter 13 (Universal source coding)</li> <li>Grunwald Chapters 1-3 (introduction, Kolmogorov complexity, universal coding)</li> </ul> <h2 id="week-6-channel-coding-and-mutual-information"><a class="header" href="#week-6-channel-coding-and-mutual-information">Week 6: Channel Coding and Mutual Information</a></h2> <ul> <li>Mackay Chapter 9 (channel capacity)</li> <li>Cover and Thomas Chapter 2 (entropy, mutual information properties)</li> </ul> <h2 id="week-7-compression-and-llm-hallucinations"><a class="header" href="#week-7-compression-and-llm-hallucinations">Week 7: Compression and LLM Hallucinations</a></h2> <ul> <li>Read paper <a href="https://arxiv.org/abs/2509.11208">Predictable Compression Failures: Why Language Models Actually Hallucinate</a></li> </ul> <h2 id="week-8-rate-distortion-and-information-bottleneck"><a class="header" href="#week-8-rate-distortion-and-information-bottleneck">Week 8: Rate Distortion and Information Bottleneck</a></h2> <ul> <li>Cover and Thomas Chapter 10 (rate distortion)</li> <li>Tishby 1999 - The information bottleneck method</li> <li>Alemi 2017 - Deep variational information bottleneck</li> </ul> <h2 id="week-9-neural-compression-optional"><a class="header" href="#week-9-neural-compression-optional">Week 9: Neural Compression (Optional)</a></h2> <ul> <li>Han 2016 - Deep Compression</li> <li>Practical LLM quantization (GPTQ, AWQ)</li> </ul> <h2 id="references-6"><a class="header" href="#references-6">References</a></h2> <ol> <li>Book: Information Theory, Inference, and Learning Algorithms - David MacKay.</li> <li>Book: Elements of Information Theory - Cover & Thomas</li> <li>Monograph: High-Dimensional Probability - Vershynin</li> <li>Paper: Predictable Compression Failures (arXiv:2509.11208).</li> <li>Book: The minimum description length principle - Grunwald.</li> </ol> </main> <nav class="nav-wrapper" aria-label="Page navigation">  <div style="clear: both"></div> </nav> </div> </div> <nav class="nav-wide-wrapper" aria-label="Page navigation"> </nav> </div> <script> window.playground_copyable = true; </script> <script src="elasticlunr.min.js"></script> <script src="mark.min.js"></script> <script src="searcher.js"></script> <script src="clipboard.min.js"></script> <script src="highlight.js"></script> <script src="book.js"></script>  <script src="mermaid.min.js"></script> <script src="mermaid-init.js"></script> <script> window.addEventListener('load', function() { window.setTimeout(window.print, 100); }); </script> </div> </body> </html>

Keyboard shortcuts

Chux's Notebook

Chux's Notebook

About Me

Current Focus

2025 Goals

Research Questions

Thoughts

On ML Experiments